From 3fbe4563125e8f89e8dbd73c5b93a6a000ba2354 Mon Sep 17 00:00:00 2001
From: Stef Smeets <s.smeets@esciencecenter.nl>
Date: Tue, 4 Jun 2024 11:56:49 +0200
Subject: [PATCH] Add method to create a graph from jumps

---
 src/gemdat/jumps.py      | 36 ++++++++++++++++++++++++++++++++++++
 src/gemdat/trajectory.py |  7 +++++++
 2 files changed, 43 insertions(+)

diff --git a/src/gemdat/jumps.py b/src/gemdat/jumps.py
index 3e4bf912..f2d4bdfd 100644
--- a/src/gemdat/jumps.py
+++ b/src/gemdat/jumps.py
@@ -5,6 +5,7 @@
 from math import ceil
 from typing import TYPE_CHECKING, Callable
 
+import networkx as nx
 import numpy as np
 import pandas as pd
 from pymatgen.core.units import FloatWithUnit
@@ -308,6 +309,41 @@ def jumps_counter(self) -> Counter:
         )
         return jumps
 
+    def activation_energy_between_sites(self, start: str, stop: str) -> float:
+        raise NotImplementedError
+
+    def to_graph(self) -> nx.DiGraph:
+        """Create a graph from jumps data.
+
+        The edges are weighted by the activation energy.
+
+        Returns
+        -------
+        G : nx.DiGraph
+            A networkx DiGraph object.
+        """
+        atom_percentage = {
+            site.label: site.species.num_atoms for site in self.transitions.occupancy()
+        }
+
+        attempt_freq, _ = self.trajectory.metrics().attempt_frequency()
+        temperature = self.trajectory.metadata['temperature']
+        kBT = Boltzmann * temperature
+
+        G = nx.DiGraph()
+
+        for (start, stop), n_jumps in self.jumps_counter().items():
+            time_perc = atom_percentage[start] * self.trajectory.total_time
+
+            eff_rate = n_jumps / time_perc
+
+            e_act = -np.log(eff_rate / attempt_freq) * kBT
+            e_act /= elementary_charge
+
+            G.add_edge(start, stop, e_act=e_act)
+
+        return G
+
     def split(self, n_parts: int) -> list[Jumps]:
         """Split the jumps into parts.
 
diff --git a/src/gemdat/trajectory.py b/src/gemdat/trajectory.py
index 37851854..6e02c3ac 100644
--- a/src/gemdat/trajectory.py
+++ b/src/gemdat/trajectory.py
@@ -19,6 +19,7 @@
 if TYPE_CHECKING:
     from pymatgen.core import Structure
 
+    from .metrics import SimulationMetrics
     from .transitions import Transitions
     from .volume import Volume
 
@@ -520,6 +521,12 @@ def transitions_between_sites(
             site_inner_fraction=site_inner_fraction,
         )
 
+    def metrics(self) -> SimulationMetrics:
+        """See [gemdat.SimulationMetrics][] for more info."""
+        from .metrics import SimulationMetrics
+
+        return SimulationMetrics(trajectory=self)
+
     def plot_displacement_per_atom(self, **kwargs):
         """See [gemdat.plots.displacement_per_atom][] for more info."""
         from gemdat import plots