Repository: PHOTOX/ABIN
Pull request #206: Python script for analyzing TSH trajectories from ABIN.
Author: login: JanosJiri
Created: 2025-05-06T17:55:31Z | Merged: 2025-05-14T21:20:26Z

The script plots PES along the trajectory, electronic populations, total and kinetic energy from either LZSH or FSSH. All is in one plot and allows analysis of a single trajectory run.


Diff:
diff --git a/.github/workflows/format-lint.yml b/.github/workflows/format-lint.yml
index c9187080..2ae20aef 100644
--- a/.github/workflows/format-lint.yml
+++ b/.github/workflows/format-lint.yml
@@ -34,3 +34,13 @@ jobs:
 
     - name: Lint check
       run: fortitude check src/
+
+  scripts:
+    name: Test and lint utililities and analysis scripts
+    runs-on: ubuntu-24.04
+    steps:
+    - uses: actions/checkout@v4
+    - uses: astral-sh/setup-uv@v6
+    - name: Run scripts
+      # TODO: Test the script on actual TSH trajectory
+      run: uv run -p 3.9 utils/analysis/tsh_traj_anal.py -h
diff --git a/utils/analyzeSH.sh b/utils/analysis/analyzeSH.sh
similarity index 100%
rename from utils/analyzeSH.sh
rename to utils/analysis/analyzeSH.sh
diff --git a/utils/analysis/tsh_traj_anal.py b/utils/analysis/tsh_traj_anal.py
new file mode 100755
index 00000000..2e035f1b
--- /dev/null
+++ b/utils/analysis/tsh_traj_anal.py
@@ -0,0 +1,181 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# # NOTE: PyQT6 dependency is a workaround for:
+# # https://github.com/astral-sh/uv/issues/6893
+# dependencies = [
+#     "matplotlib",
+#     "numpy",
+#     "PyQT6",
+# ]
+# ///
+"""Analysis script for Surface Hopping trajectories generated by ABIN
+
+Script for analyzing potential energy surfaces, energies and electronic populations along
+individual trajectories from ABIN's Fewest Switches or Landau-Zener surface hopping simulations.
+
+The script should be launched within the ABIN trajectory folder as:
+
+    python3 tsh_traj_analysis.py
+
+or:
+
+    ./tsh_traj_analysis.py
+
+
+You can also launch the script via 'uv' (https://github.com/astral-sh/uv)
+which will automatically install all needed dependencies (numpy, matplotlib)
+in a separate virtual environment:
+
+    uv run tsh_traj_analysis.py
+
+Author: Jiri Janos 2025
+"""
+
+import os.path
+
+# ABIN output files that we need
+POPULATION_FILE = "pop.dat"
+PES_FILE = "PES.dat"
+ENERGY_FILE = "energies.dat"
+
+AUtoEV = 27.2114
+
+
+def parse_cmd():
+    """Parse command-line arguments"""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Analyze Surface Hopping trajectories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("-s", "--save-fig", action="store_true", help="Save png figure")
+    parser.add_argument("-eu", "--energy-units", choices=["eV", "au"], default="eV", help="Energy units")
+    parser.add_argument(
+        "-es", 
+        "--energy-shift",
+        type=float,
+        default=None,
+        help="Energy shift parameter in energy-units. All energy curves will be shifted by this energy shift parameter. If the option is not specified, the ground-state energy at time zero will be set to 0.",
+    )
+    parser.add_argument("-n", "--nstates", type=int, help="Number of states to plot (default: all)")
+    return parser.parse_args()
+
+
+config = parse_cmd()
+
+for fname in (POPULATION_FILE, PES_FILE, ENERGY_FILE):
+    if not os.path.exists(fname):
+        exit(f"ERROR: file '{fname}' does not exist")
+
+# Lazy imports to speed up CLI help printing
+import numpy as np  # noqa: E402
+import matplotlib.pyplot as plt  # noqa: E402
+
+# reading data from files
+pop = np.genfromtxt(POPULATION_FILE)
+data = np.genfromtxt(PES_FILE)
+energies = np.genfromtxt(ENERGY_FILE)
+
+# Set nstates to all unless specified otherwise, or check if nstates is valid
+nstates_all = len(data.T) - 1
+nstates = config.nstates
+if nstates is None:
+    nstates = nstates_all
+elif nstates > nstates_all:
+    exit(f"ERROR: nstates ({nstates}) is larger than the number of states in the data ({nstates_all})")
+elif nstates < 1:
+    exit(f"ERROR: nstates ({nstates}) must be at least 1")
+
+# convert energy units
+if config.energy_units == "eV":
+    data.T[1:, :] = data.T[1:, :] * AUtoEV
+    energies.T[1:, :] = energies.T[1:, :] * AUtoEV
+    enunits = "eV"
+else:
+    enunits = "a.u."
+
+# Shift energies
+if config.energy_shift is None:
+    minE = -np.min(data.T[1:, 0])
+else:
+    minE = config.energy_shift
+data.T[1:, :] = data.T[1:, :] + minE # shifting PES
+energies.T[1, :] = energies.T[1, :] + minE # shifting kinetic energy
+energies.T[-1, :] = energies.T[-1, :] + minE # shifting total energy
+
+### plotting
+colors = plt.cm.viridis(np.linspace(0, 0.8, nstates))
+fig, axs = plt.subplots(4, 1, figsize=(8, 7.5), gridspec_kw={"height_ratios": [1, 2, 1, 1]}, sharex=True)
+
+# plotting populations
+for i in range(nstates):
+    axs[0].plot(pop.T[0, :], pop.T[i + 2, :], color=colors[i], label=rf"$S_{i}$")
+
+axs[0].plot(
+    pop.T[0, :],
+    pop.T[1, :] - 1,
+    color="black",
+    linestyle="dashed",
+    label="act. state\nindex",
+)
+
+axs[0].set_ylabel("El. populations")
+axs[0].legend(labelspacing=0)
+
+# plotting PES
+for i in range(0, nstates):
+    axs[1].plot(data.T[0, :], data.T[i + 1, :], color=colors[i], label=rf"$S_{i}$")
+
+axs[1].scatter(energies.T[0, 0:-1:5], energies.T[1, 0:-1:5], s=5, alpha=0.5, color="black")
+axs[1].plot(
+    energies.T[0, :],
+    energies.T[1, :],
+    linewidth=8,
+    alpha=0.2,
+    color="black",
+    label="active",
+)
+axs[1].plot(
+    energies.T[0, :],
+    energies.T[-1, :],
+    color="black",
+    alpha=0.5,
+    linestyle="dotted",
+    label=r"$E_{tot}$",
+)
+
+axs[1].set_ylabel(f"Energy ({enunits})")
+axs[1].legend(labelspacing=0)
+
+# plotting kinetic energy
+axs[2].plot(energies.T[0, :], energies.T[2, :], color="black", alpha=1, label=r"$E_k$")
+axs[2].set_ylabel(f"Kin. energy ({enunits})")
+
+# plotting total energy
+for i in range(0, nstates):
+    axs[3].scatter(
+        energies.T[0, pop.T[1, :] == (i + 1)],
+        energies.T[-1, pop.T[1, :] == (i + 1)] - energies.T[-1, 0],
+        color=colors[i],
+        s=10,
+    )
+axs[3].plot(energies.T[0, :], energies.T[-1, :] - energies.T[-1, 0], color="black", alpha=0.5)
+axs[3].plot([energies.T[0, 0], energies.T[0, -1]], [0, 0], color="black", linestyle="dashed")
+
+axs[3].set_ylabel(f"Total energy ({enunits})")
+axs[3].set_xlabel("Time (fs)")
+
+# setting minor ticks
+for axis in axs:
+    axis.minorticks_on()
+    axis.tick_params("both", direction="in", which="both", top=True, right=True)
+
+plt.tight_layout()
+plt.subplots_adjust(hspace=0)
+
+# save figure
+if config.save_fig:
+    plt.savefig("PES_pop", dpi=300)
+plt.show()


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR introduces a Python script for analyzing TSH trajectories from ABIN simulations by plotting potential energy surfaces, electronic populations, and energy data from either LZSH or FSSH runs.  
- Adds a command-line interface for specifying input files and options for energy unit conversion and curve shifting.  
- Implements file existence checks and plotting of trajectories and energy profiles based on generated data.




state: COMMENTED
author_association: NONE
submitted_at: 2025-05-06T17:58:15Z
user: login: danielhollas
body: This is great, thanks! I made some insignificant changes and added inline metadata so that the script can be run via `uv run`. I have some questions / suggestions about the command line arguments.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-12T05:12:34Z
user: login: JanosJiri
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-12T06:39:00Z
user: login: JanosJiri
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-12T06:58:20Z
user: login: danielhollas
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T12:21:47Z
user: login: danielhollas
body: Thank you for the changes! I have two questions about the energy shifting, but otherwise I am happy, this will be very useful!

I don't think we need a new top-level `analysis` directory at this point, so I moved it to `utils/`. At some point we might want to organize the `utils/` directory a bit better, but not here. 
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T12:49:17Z
user: login: JanosJiri
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T15:11:34Z
user: login: JanosJiri
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T15:13:15Z
user: login: danielhollas
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T21:20:00Z

Review comments:
diff_hunk: @@ -0,0 +1,171 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# dependencies = [
+#     "matplotlib",
+#     "numpy",
+#     "PyQT6",
+# ]
+# ///
+"""Analysis script for Surface Hopping trajectories generated by ABIN
+
+Script for analyzing potential energy surfaces, energies and electronic populations along
+individual trajectories from ABIN's Fewest Switches or Landau-Zener surface hopping simulations.
+
+The script should be launched within the ABIN trajectory folder as:
+
+    python3 tsh_traj_analysis.py
+
+or:
+
+    ./tsh_traj_analysis.py
+
+
+Author: Jiri Janos 2025
+"""
+
+def parse_cmd():
+    """Parse command-line arguments"""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Analyze Surface Hopping trajectories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("-p", "--print", action="store_true", help="Save png figure")
path: analysis/tsh_traj_anal.py
user: login: danielhollas
body: ```suggestion
    parser.add_argument("-s", "--save-fig", action="store_true", help="Save png figure")
```
created_at: 2025-05-11T20:04:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,171 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# dependencies = [
+#     "matplotlib",
+#     "numpy",
+#     "PyQT6",
+# ]
+# ///
+"""Analysis script for Surface Hopping trajectories generated by ABIN
+
+Script for analyzing potential energy surfaces, energies and electronic populations along
+individual trajectories from ABIN's Fewest Switches or Landau-Zener surface hopping simulations.
+
+The script should be launched within the ABIN trajectory folder as:
+
+    python3 tsh_traj_analysis.py
+
+or:
+
+    ./tsh_traj_analysis.py
+
+
+Author: Jiri Janos 2025
+"""
+
+def parse_cmd():
+    """Parse command-line arguments"""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Analyze Surface Hopping trajectories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("-p", "--print", action="store_true", help="Save png figure")
+    parser.add_argument("-c", "--convert", action="store_false", help="Convert a.u. to eV")
path: analysis/tsh_traj_anal.py
user: login: danielhollas
body: I would instead do the same thing we have in promdens
```suggestion
    parser.add_argument("--energy-units", choices=["eV", "au"], default="eV", help="Energy units")
```
created_at: 2025-05-11T20:06:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,171 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# dependencies = [
+#     "matplotlib",
+#     "numpy",
+#     "PyQT6",
+# ]
+# ///
+"""Analysis script for Surface Hopping trajectories generated by ABIN
+
+Script for analyzing potential energy surfaces, energies and electronic populations along
+individual trajectories from ABIN's Fewest Switches or Landau-Zener surface hopping simulations.
+
+The script should be launched within the ABIN trajectory folder as:
+
+    python3 tsh_traj_analysis.py
+
+or:
+
+    ./tsh_traj_analysis.py
+
+
+Author: Jiri Janos 2025
+"""
+
+def parse_cmd():
+    """Parse command-line arguments"""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Analyze Surface Hopping trajectories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("-p", "--print", action="store_true", help="Save png figure")
+    parser.add_argument("-c", "--convert", action="store_false", help="Convert a.u. to eV")
+    parser.add_argument(
+        "-s",
+        "--shift",
+        action="store_false",
+        help="Shift curves so that the lowest state minimum has 0 energy",
+    )
+    parser.add_argument("-n", "--nstates", type=int, help="Number of states to plot (default: all)")
+    parser.add_argument(
+        "-pop",
+        "--population",
path: analysis/tsh_traj_anal.py
user: login: danielhollas
body: Is there an actual use for these arguments? Why would the files be named differently?
created_at: 2025-05-11T20:07:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,171 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# dependencies = [
path: analysis/tsh_traj_anal.py
user: login: danielhollas
body: Added via

```
uv add --script tsh_traj_anal.py numpy matplotlib PyQT6
```
created_at: 2025-05-11T20:08:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,171 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# dependencies = [
+#     "matplotlib",
+#     "numpy",
+#     "PyQT6",
+# ]
+# ///
+"""Analysis script for Surface Hopping trajectories generated by ABIN
+
+Script for analyzing potential energy surfaces, energies and electronic populations along
+individual trajectories from ABIN's Fewest Switches or Landau-Zener surface hopping simulations.
+
+The script should be launched within the ABIN trajectory folder as:
+
+    python3 tsh_traj_analysis.py
+
+or:
+
+    ./tsh_traj_analysis.py
+
+
+Author: Jiri Janos 2025
+"""
+
+def parse_cmd():
+    """Parse command-line arguments"""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Analyze Surface Hopping trajectories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("-p", "--print", action="store_true", help="Save png figure")
+    parser.add_argument("-c", "--convert", action="store_false", help="Convert a.u. to eV")
+    parser.add_argument(
+        "-s",
+        "--shift",
path: analysis/tsh_traj_anal.py
user: login: danielhollas
body: Maybe
```suggestion
        "--energy-shift",
```
created_at: 2025-05-12T05:11:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,171 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# dependencies = [
path: analysis/tsh_traj_anal.py
user: login: JanosJiri
body: Just a question, why is PyQT6 needed here?
created_at: 2025-05-12T06:38:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,171 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# dependencies = [
+#     "matplotlib",
+#     "numpy",
+#     "PyQT6",
+# ]
+# ///
+"""Analysis script for Surface Hopping trajectories generated by ABIN
+
+Script for analyzing potential energy surfaces, energies and electronic populations along
+individual trajectories from ABIN's Fewest Switches or Landau-Zener surface hopping simulations.
+
+The script should be launched within the ABIN trajectory folder as:
+
+    python3 tsh_traj_analysis.py
+
+or:
+
+    ./tsh_traj_analysis.py
+
+
+Author: Jiri Janos 2025
+"""
+
+def parse_cmd():
+    """Parse command-line arguments"""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Analyze Surface Hopping trajectories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("-p", "--print", action="store_true", help="Save png figure")
+    parser.add_argument("-c", "--convert", action="store_false", help="Convert a.u. to eV")
+    parser.add_argument(
+        "-s",
+        "--shift",
+        action="store_false",
+        help="Shift curves so that the lowest state minimum has 0 energy",
+    )
+    parser.add_argument("-n", "--nstates", type=int, help="Number of states to plot (default: all)")
+    parser.add_argument(
+        "-pop",
+        "--population",
path: analysis/tsh_traj_anal.py
user: login: JanosJiri
body: Originally, the script was not intended only for ABIN, but I never used it for anything else. I'll remove it.
created_at: 2025-05-12T06:57:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,171 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# dependencies = [
path: analysis/tsh_traj_anal.py
user: login: danielhollas
body: Good question. When using python installations from uv, the `plt.show` command did not work for me (no graph was interactively displayed). This is a known issue: https://github.com/astral-sh/uv/issues/6893

So this is a workaround for that. It's not ideal, but I think better then potentially failing script. I'll add a comment about it.
created_at: 2025-05-13T12:21:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,178 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# # NOTE: PyQT6 dependency is a workaround for:
+# # https://github.com/astral-sh/uv/issues/6893
+# dependencies = [
+#     "matplotlib",
+#     "numpy",
+#     "PyQT6",
+# ]
+# ///
+"""Analysis script for Surface Hopping trajectories generated by ABIN
+
+Script for analyzing potential energy surfaces, energies and electronic populations along
+individual trajectories from ABIN's Fewest Switches or Landau-Zener surface hopping simulations.
+
+The script should be launched within the ABIN trajectory folder as:
+
+    python3 tsh_traj_analysis.py
+
+or:
+
+    ./tsh_traj_analysis.py
+
+
+You can also launch the script via 'uv' (https://github.com/astral-sh/uv)
+which will automatically install all needed dependencies (numpy, matplotlib)
+in a separate virtual environment:
+
+    uv run tsh_traj_analysis.py
+
+Author: Jiri Janos 2025
+"""
+
+import os.path
+
+# ABIN output files that we need
+POPULATION_FILE = "pop.dat"
+PES_FILE = "PES.dat"
+ENERGY_FILE = "energies.dat"
+
+AUtoEV = 27.2114
+
+
+def parse_cmd():
+    """Parse command-line arguments"""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Analyze Surface Hopping trajectories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("-s", "--save-fig", action="store_true", help="Save png figure")
+    parser.add_argument("-eu", "--energy-units", choices=["eV", "au"], default="eV", help="Energy units")
+    parser.add_argument(
+        "-es",
+        "--energy-shift",
+        action="store_false",
+        help="Shift curves so that the lowest state minimum has 0 energy",
+    )
path: utils/analysis/tsh_traj_anal.py
user: login: danielhollas
body: This option is currently confusing to me. If I understand correctly, specifying this option actually **turns off** the energy shift?

What if we instead have an option that specifies the zero energy (probably in atomic units)? For example

```suggestion
    parser.add_argument(
        "-ze",
        "--zero-energy",
        type=float
        default=None,
        help="Specify zero energy (in atomic units). All energy curves will be shifted with respect to this energy.",
    )
```

To turn off energy shifting, the user then could pass (`--energy-shift 0`).
And this would also allow to pass any value, which would be handy if you were analyzing multiple trajectories and wanted to have consistent E0)
created_at: 2025-05-13T12:45:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,178 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# # NOTE: PyQT6 dependency is a workaround for:
+# # https://github.com/astral-sh/uv/issues/6893
+# dependencies = [
+#     "matplotlib",
+#     "numpy",
+#     "PyQT6",
+# ]
+# ///
+"""Analysis script for Surface Hopping trajectories generated by ABIN
+
+Script for analyzing potential energy surfaces, energies and electronic populations along
+individual trajectories from ABIN's Fewest Switches or Landau-Zener surface hopping simulations.
+
+The script should be launched within the ABIN trajectory folder as:
+
+    python3 tsh_traj_analysis.py
+
+or:
+
+    ./tsh_traj_analysis.py
+
+
+You can also launch the script via 'uv' (https://github.com/astral-sh/uv)
+which will automatically install all needed dependencies (numpy, matplotlib)
+in a separate virtual environment:
+
+    uv run tsh_traj_analysis.py
+
+Author: Jiri Janos 2025
+"""
+
+import os.path
+
+# ABIN output files that we need
+POPULATION_FILE = "pop.dat"
+PES_FILE = "PES.dat"
+ENERGY_FILE = "energies.dat"
+
+AUtoEV = 27.2114
+
+
+def parse_cmd():
+    """Parse command-line arguments"""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Analyze Surface Hopping trajectories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("-s", "--save-fig", action="store_true", help="Save png figure")
+    parser.add_argument("-eu", "--energy-units", choices=["eV", "au"], default="eV", help="Energy units")
+    parser.add_argument(
+        "-es",
+        "--energy-shift",
+        action="store_false",
+        help="Shift curves so that the lowest state minimum has 0 energy",
+    )
+    parser.add_argument("-n", "--nstates", type=int, help="Number of states to plot (default: all)")
+    return parser.parse_args()
+
+
+config = parse_cmd()
+
+for fname in (POPULATION_FILE, PES_FILE, ENERGY_FILE):
+    if not os.path.exists(fname):
+        exit(f"ERROR: file '{fname}' does not exist")
+
+# Lazy imports to speed up CLI help printing
+import numpy as np  # noqa: E402
+import matplotlib.pyplot as plt  # noqa: E402
+
+# reading data from files
+pop = np.genfromtxt(POPULATION_FILE)
+data = np.genfromtxt(PES_FILE)
+energies = np.genfromtxt(ENERGY_FILE)
+
+# Set nstates to all unless specified otherwise, or check if nstates is valid
+nstates_all = len(data.T) - 1
+nstates = config.nstates
+if nstates is None:
+    nstates = nstates_all
+elif nstates > nstates_all:
+    exit(f"ERROR: nstates ({nstates}) is larger than the number of states in the data ({nstates_all})")
+elif nstates < 1:
+    exit(f"ERROR: nstates ({nstates}) must be at least 1")
+
+# convert energy units
+if config.energy_units == "eV":
+    data.T[1:, :] = data.T[1:, :] * AUtoEV
+    energies.T[1:, :] = energies.T[1:, :] * AUtoEV
+    enunits = "eV"
+else:
+    enunits = "a.u."
+
+# Shift energies
+if config.energy_shift:
+    minE = np.min(data.T[1:, :])
path: utils/analysis/tsh_traj_anal.py
user: login: danielhollas
body: I did not notice this before, but this looks like the E0 is taken as the minimum S0 energy across the whole trajectory? To me this is a non-obvious choice. I would expect to simply shift the energy based on S0 at t0. What do you think?
created_at: 2025-05-13T12:47:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,178 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# # NOTE: PyQT6 dependency is a workaround for:
+# # https://github.com/astral-sh/uv/issues/6893
+# dependencies = [
+#     "matplotlib",
+#     "numpy",
+#     "PyQT6",
+# ]
+# ///
+"""Analysis script for Surface Hopping trajectories generated by ABIN
+
+Script for analyzing potential energy surfaces, energies and electronic populations along
+individual trajectories from ABIN's Fewest Switches or Landau-Zener surface hopping simulations.
+
+The script should be launched within the ABIN trajectory folder as:
+
+    python3 tsh_traj_analysis.py
+
+or:
+
+    ./tsh_traj_analysis.py
+
+
+You can also launch the script via 'uv' (https://github.com/astral-sh/uv)
+which will automatically install all needed dependencies (numpy, matplotlib)
+in a separate virtual environment:
+
+    uv run tsh_traj_analysis.py
+
+Author: Jiri Janos 2025
+"""
+
+import os.path
+
+# ABIN output files that we need
+POPULATION_FILE = "pop.dat"
+PES_FILE = "PES.dat"
+ENERGY_FILE = "energies.dat"
+
+AUtoEV = 27.2114
+
+
+def parse_cmd():
+    """Parse command-line arguments"""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Analyze Surface Hopping trajectories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("-s", "--save-fig", action="store_true", help="Save png figure")
+    parser.add_argument("-eu", "--energy-units", choices=["eV", "au"], default="eV", help="Energy units")
+    parser.add_argument(
+        "-es",
+        "--energy-shift",
+        action="store_false",
+        help="Shift curves so that the lowest state minimum has 0 energy",
+    )
path: utils/analysis/tsh_traj_anal.py
user: login: JanosJiri
body: That's a good point. I added it ages ago and probably used it once in my life. Your suggestion makes complete sense.

created_at: 2025-05-14T15:11:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,178 @@
+#!/bin/env python3
+# /// script
+# requires-python = ">=3.9"
+# # NOTE: PyQT6 dependency is a workaround for:
+# # https://github.com/astral-sh/uv/issues/6893
+# dependencies = [
+#     "matplotlib",
+#     "numpy",
+#     "PyQT6",
+# ]
+# ///
+"""Analysis script for Surface Hopping trajectories generated by ABIN
+
+Script for analyzing potential energy surfaces, energies and electronic populations along
+individual trajectories from ABIN's Fewest Switches or Landau-Zener surface hopping simulations.
+
+The script should be launched within the ABIN trajectory folder as:
+
+    python3 tsh_traj_analysis.py
+
+or:
+
+    ./tsh_traj_analysis.py
+
+
+You can also launch the script via 'uv' (https://github.com/astral-sh/uv)
+which will automatically install all needed dependencies (numpy, matplotlib)
+in a separate virtual environment:
+
+    uv run tsh_traj_analysis.py
+
+Author: Jiri Janos 2025
+"""
+
+import os.path
+
+# ABIN output files that we need
+POPULATION_FILE = "pop.dat"
+PES_FILE = "PES.dat"
+ENERGY_FILE = "energies.dat"
+
+AUtoEV = 27.2114
+
+
+def parse_cmd():
+    """Parse command-line arguments"""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Analyze Surface Hopping trajectories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("-s", "--save-fig", action="store_true", help="Save png figure")
+    parser.add_argument("-eu", "--energy-units", choices=["eV", "au"], default="eV", help="Energy units")
+    parser.add_argument(
+        "-es",
+        "--energy-shift",
+        action="store_false",
+        help="Shift curves so that the lowest state minimum has 0 energy",
+    )
+    parser.add_argument("-n", "--nstates", type=int, help="Number of states to plot (default: all)")
+    return parser.parse_args()
+
+
+config = parse_cmd()
+
+for fname in (POPULATION_FILE, PES_FILE, ENERGY_FILE):
+    if not os.path.exists(fname):
+        exit(f"ERROR: file '{fname}' does not exist")
+
+# Lazy imports to speed up CLI help printing
+import numpy as np  # noqa: E402
+import matplotlib.pyplot as plt  # noqa: E402
+
+# reading data from files
+pop = np.genfromtxt(POPULATION_FILE)
+data = np.genfromtxt(PES_FILE)
+energies = np.genfromtxt(ENERGY_FILE)
+
+# Set nstates to all unless specified otherwise, or check if nstates is valid
+nstates_all = len(data.T) - 1
+nstates = config.nstates
+if nstates is None:
+    nstates = nstates_all
+elif nstates > nstates_all:
+    exit(f"ERROR: nstates ({nstates}) is larger than the number of states in the data ({nstates_all})")
+elif nstates < 1:
+    exit(f"ERROR: nstates ({nstates}) must be at least 1")
+
+# convert energy units
+if config.energy_units == "eV":
+    data.T[1:, :] = data.T[1:, :] * AUtoEV
+    energies.T[1:, :] = energies.T[1:, :] * AUtoEV
+    enunits = "eV"
+else:
+    enunits = "a.u."
+
+# Shift energies
+if config.energy_shift:
+    minE = np.min(data.T[1:, :])
path: utils/analysis/tsh_traj_anal.py
user: login: JanosJiri
body: Yep, that's also a good point. I would modify it to the minimum of PES at t0. If we have triplets in LZ, the lowest triplet could be below S0, e.g., for oxygen. But otherwise it should be the same as taking S0 energy at t0.
created_at: 2025-05-14T15:13:15Z
author_association: CONTRIBUTOR

PR comments:
user: login: JanosJiri
created_at: 2025-05-06T17:57:49Z
body: Hi @danielhollas, I've created a Python script which I use to analyze TSH calculations in ABIN. I thought it might be handy to attach it to ABIN as people could use it for analysis. Please, have a look and let me know if you fancy having such a script with ABIN. 
author_association: CONTRIBUTOR
user: login: codecov[bot]
created_at: 2025-05-06T18:01:48Z
body: ## [Codecov](https://app.codecov.io/gh/PHOTOX/ABIN/pull/206?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PHOTOX) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 84.05%. Comparing base [(`c83c105`)](https://app.codecov.io/gh/PHOTOX/ABIN/commit/c83c105075543773e7334135c0c8bacd92080664?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PHOTOX) to head [(`0d3ff08`)](https://app.codecov.io/gh/PHOTOX/ABIN/commit/0d3ff08f23c98be4cd91c8d0aee853f17f536e2b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PHOTOX).
> Report is 1 commits behind head on master.

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://app.codecov.io/gh/PHOTOX/ABIN/pull/206/graphs/tree.svg?width=650&height=150&src=pr&token=YU1323XAT1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PHOTOX)](https://app.codecov.io/gh/PHOTOX/ABIN/pull/206?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PHOTOX)

```diff
@@           Coverage Diff           @@
##           master     #206   +/-   ##
=======================================
  Coverage   84.05%   84.05%           
=======================================
  Files          47       47           
  Lines        6599     6599           
  Branches      764      764           
=======================================
  Hits         5547     5547           
  Misses        853      853           
  Partials      199      199           
```

| [Flag](https://app.codecov.io/gh/PHOTOX/ABIN/pull/206/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PHOTOX) | Coverage Δ | |
|---|---|---|
| [unittests](https://app.codecov.io/gh/PHOTOX/ABIN/pull/206/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PHOTOX) | `18.41% <ø> (ø)` | |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PHOTOX#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>
<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: danielhollas
created_at: 2025-05-09T09:39:06Z
body: > I thought it might be handy to attach it to ABIN as people could use it for analysis. Please, have a look and let me know if you fancy having such a script with ABIN.

Yes please! I'll have a closer look later but this looks very handy of course! :clap: :heart_eyes: 

(I might do a quick pass over the script if you don't mind. But let me know if you prefer to iterate over review)
author_association: CONTRIBUTOR
user: login: JanosJiri
created_at: 2025-05-11T14:20:56Z
body: Go for it, no need to iterate over review.
author_association: CONTRIBUTOR
user: login: JanosJiri
created_at: 2025-05-12T07:26:33Z
body: I modified the script according to your comments plus I added some checks for the nstates variable. Please, have a look. 
author_association: CONTRIBUTOR
user: login: JanosJiri
created_at: 2025-05-14T16:29:18Z
body: I worked on the suggestions regarding the `energy-shift` parameter. In the end, I combined your suggestion with the original version. Currently, if you don't use `--energy-shift`, all the energies are shifted such that you start with the lowest state at zero energy.  If you specify some value for energy shift, it will shift the original energies by the value specified. I kept it that way because if I want to look at multiple trajectories and have them all starting at the same position, I use `--energy-shift`. If I want to check the absolute energies, I use `--energy-shift 0`. And if I want to shift just one of the trajectories to see how it compares to others, I would use `--energy-shift xx` for the one. That's all the regimes I ever wanted to use. Let me know what you think about it. 
author_association: CONTRIBUTOR
user: login: danielhollas
created_at: 2025-05-14T21:19:50Z
body: Perfect! That's actually what I had in mind, but didn't write out clearly apparently. Thanks for bearing with me. :-)
author_association: CONTRIBUTOR