Added Jumps class

src/gemdat/collective.py

@@ -3,11 +3,15 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, Tuple
+
 import numpy as np
 from pymatgen.core import Lattice, Structure
 
 from .caching import weak_lru_cache
-from .transitions import Transitions
+
+if TYPE_CHECKING:
+    from .jumps import Jumps
 
 
 class Collective:
@@ -24,7 +28,7 @@ class Collective:
     """
 
     def __init__(self,
-                 transitions: Transitions,
+                 jumps: Jumps,
                  structure: Structure,
                  lattice: Lattice,
                  max_steps: int,
@@ -33,8 +37,8 @@ def __init__(self,
 
         Parameters
         ----------
-        transitions : Transitions
-            Input transition events
+        jumps : Jumps
+            Input jump events
         structure : pymatgen.core.structure.Structure
             Structure with list of jump sites
         lattice : pymatgen.core.lattice.Lattice
@@ -44,7 +48,7 @@ def __init__(self,
         max_dist : float
             Maximum distance for collective motions in Angstrom
         """
-        self.transitions = transitions
+        self.jumps = jumps
         self.structure = structure
         self.lattice = lattice
         self.max_steps = max_steps
@@ -59,51 +63,52 @@ def _compute(self):
         max_steps = self.max_steps
         max_dist = self.max_dist
 
-        time_col = 4
-        events = self.transitions.events
-        event_order = events[:, time_col].argsort()
+        events = self.jumps.as_dataframe()
+        events = events.sort_values(['stop time', 'start time'],
+                                    ignore_index=True)
 
         coll_jumps = []
         collective = []
+        collective_matrix = np.full((len(events), len(events)), False)
 
-        n_solo_jumps = 0
-
-        for i, event_i in enumerate(event_order[:-1]):
-
-            for j, event_j in enumerate(event_order[i + 1:], start=1):
-
-                atom_i, start_site_i, stop_site_i, _, stop_time_i = events[
-                    event_i]
-                atom_j, start_site_j, stop_site_j, _, stop_time_j = events[
-                    event_j]
-
-                if stop_time_j - stop_time_i > max_steps:
+        # Compare all pairs
+        for i, event_i in events[:-1].iterrows():
+            for j, event_j in events[i + 1:].iterrows():
+                if event_j['start time'] - event_i['stop time'] > max_steps:
                     break
-
-                if atom_i == atom_j:
-                    n_solo_jumps += 1
+                if event_i['start time'] - event_j['stop time'] > max_steps:
+                    continue
+                if event_i['atom index'] == event_j['atom index']:
                     continue
 
-                a = structure.frac_coords[[start_site_i, stop_site_i]]
-                b = structure.frac_coords[[start_site_j, stop_site_j]]
+                a = structure.frac_coords[[
+                    event_i['start site'], event_i['destination site']
+                ]]
+                b = structure.frac_coords[[
+                    event_j['start site'], event_j['destination site']
+                ]]
 
                 dists = lattice.get_all_distances(a, b)
 
                 if np.any(dists < max_dist):
                     collective.append((event_i, event_j))
-                    coll_jumps.append(((start_site_i, stop_site_i),
-                                       (start_site_j, stop_site_j)))
+                    coll_jumps.append(
+                        ((event_i['start site'], event_i['destination site']),
+                         (event_j['start site'], event_j['destination site'])))
+                    collective_matrix[i, j] = True
+                    collective_matrix[j, i] = True
 
-                else:
-                    n_solo_jumps += 1
+        # Get solo jumps from the collective matrix
+        self.n_solo_jumps = len(events) - np.any(collective_matrix,
+                                                 axis=0).sum()
 
         self.collective = collective
         self.coll_jumps = coll_jumps
-        self.n_solo_jumps = n_solo_jumps
+        self.collective_matrix = collective_matrix
 
     @weak_lru_cache()
     def matrix(self) -> np.ndarray:
-        """Calculate collective jumps matrix.
+        """Collective jumps matrix.
 
         Returns
         -------
@@ -134,19 +139,26 @@ def matrix(self) -> np.ndarray:
         return collective_matrix
 
     @weak_lru_cache()
-    def multiple_collective(self) -> np.ndarray:
+    def multiple_collective(self) -> Tuple[np.ndarray, np.ndarray]:
         """Find jumps that occur collectively multiple times.
 
-        Only returns non-unique jumps
+        returns collective jumps and their occurence
 
         Returns
         -------
-        multiple_collective : np.ndarray
-            Dictionary with indices of sites between which jumps happen and their counts.
+        multiple_collective : np.ndarray, np.ndarray
+            Result of a np.unique on the collective jump pairs
         """
-        collective = self.collective
-        coll_sorted = np.sort(np.array(collective).flatten())
-        difference = np.diff(coll_sorted, prepend=0)
-        multiple_collective = coll_sorted[difference == 0]
-
-        return multiple_collective
+        collective = np.array(dtype=[('start', int), ('stop', int)],
+                              object=[[(event_i['start site'],
+                                        event_i['destination site']),
+                                       (event_j['start site'],
+                                        event_j['destination site'])]
+                                      for event_i, event_j in self.collective])
+        # Sort jumps so equal jumps are orderd similarily
+        collective = np.sort(collective, axis=1)
+
+        # Counts how often collective jumps occur
+        jumps, counts = np.unique(collective, axis=0, return_counts=True)
+
+        return jumps, counts