Return structure with occupancy set as Composition

src/gemdat/jumps.py

@@ -227,13 +227,15 @@ def collective(self, max_dist: float = 1) -> Collective:
 
     @weak_lru_cache()
     def activation_energies(
-            self) -> dict[tuple[str, str], tuple[float, float]]:
+            self, n_parts: int) -> dict[tuple[str, str], tuple[float, float]]:
         """Calculate activation energies for jumps (UNITS?).
 
         Returns
         -------
         e_act : dict[tuple[str, str], tuple[float, float]]
             Dictionary with jump activation energies and standard deviations between site pairs.
+        n_parts : int
+            Number of parts to split transitions/jumps into for statistics
         """
         trajectory = self.transitions.trajectory
         attempt_freq, _ = SimulationMetrics(trajectory).attempt_frequency()
@@ -242,16 +244,18 @@ def activation_energies(
 
         temperature = trajectory.metadata['temperature']
 
-        atom_locations_parts = self.sites.atom_locations_parts(
-            self.transitions)
-        parts = self.jumps_counter_parts(self.sites.n_parts)
+        atom_locations_parts = [
+            self.sites.atom_locations(part)
+            for part in self.transitions.split(n_parts)
+        ]
+        parts = [part.jumps_counter() for part in self.split(n_parts)]
 
         for site_pair in self.sites.site_pairs:
             site_start, site_stop = site_pair
 
             n_jumps = np.array([part[site_pair] for part in parts])
 
-            part_time = trajectory.total_time / self.sites.n_parts
+            part_time = trajectory.total_time / n_parts
 
             atom_percentage = np.array(
                 [part[site_start] for part in atom_locations_parts])
@@ -300,28 +304,25 @@ def split(self, n_parts) -> list[Jumps]:
         ]
 
     @weak_lru_cache()
-    def jumps_counter_parts(self, n_parts) -> list[Counter]:
-        """Return [jump counters][gemdat.sites.SitesData.jumps] per part."""
-
-        return [part.jumps_counter() for part in self.split(n_parts)]
-
-    @weak_lru_cache()
-    def rates(self) -> dict[tuple[str, str], tuple[float, float]]:
+    def rates(self,
+              n_parts: int) -> dict[tuple[str, str], tuple[float, float]]:
         """Calculate jump rates (total jumps / second).
 
         Returns
         -------
         rates : dict[tuple[str, str], tuple[float, float]]
             Dictionary with jump rates and standard deviations between site pairs
+        n_parts : int
+            Number of parts to split jumps into for statistics
         """
         rates: dict[tuple[str, str], tuple[float, float]] = {}
 
-        parts = self.jumps_counter_parts(self.sites.n_parts)
+        parts = [part.jumps_counter() for part in self.split(n_parts)]
 
         for site_pair in self.sites.site_pairs:
             n_jumps = [part[site_pair] for part in parts]
 
-            part_time = self.transitions.trajectory.total_time / self.sites.n_parts
+            part_time = self.transitions.trajectory.total_time / n_parts
             denom = self.n_floating * part_time
 
             jump_freq_mean = np.mean(n_jumps) / denom

src/gemdat/legacy.py

@@ -92,7 +92,6 @@ def analyse_md(
         structure=sites_structure,
         trajectory=trajectory,
         floating_specie=diff_elem,
-        n_parts=nr_parts,
     )
 
     transitions = Transitions.from_trajectory(trajectory=trajectory,

src/gemdat/sites.py

@@ -9,14 +9,12 @@
 import numpy as np
 from pymatgen.core import Structure
 
-from .caching import weak_lru_cache
 from .simulation_metrics import SimulationMetrics
-from .transitions import Transitions
 from .utils import is_lattice_similar
 
 if typing.TYPE_CHECKING:
-
     from gemdat.trajectory import Trajectory
+    from gemdat.transitions import Transitions
 
 NOSITE = -1
 
@@ -29,7 +27,6 @@ def __init__(
         structure: Structure,
         trajectory: Trajectory,
         floating_specie: str,
-        n_parts: int = 10,
         site_radius: Optional[float] = None,
     ):
         """Contain sites and jumps data.
@@ -42,8 +39,6 @@ def __init__(
             Input trajectory
         floating_specie : str
             Name of the floating or diffusing specie
-        n_parts : int, optional
-            Number of parts to divide transitions into for statistics
         site_radius: Optional[float]
             if set it fixes the site_radius instead of determining it
             dynamically
@@ -52,8 +47,6 @@ def __init__(
             raise ValueError(
                 'Trajectory must have constant lattice for site analysis.')
 
-        self.n_parts = n_parts
-
         self.floating_specie = floating_specie
         self.structure = structure
 
@@ -101,54 +94,7 @@ def site_pairs(self) -> list[tuple[str, str]]:
         site_pairs = product(labels, repeat=2)
         return [pair for pair in site_pairs]  # type: ignore
 
-    def occupancy_parts(self,
-                        transitions: Transitions) -> list[dict[int, int]]:
-        """Return [occupancy arrays][gemdat.transitions.Transitions.occupancy]
-        from parts."""
-        return [part.occupancy() for part in transitions.split(self.n_parts)]
-
-    def site_occupancy_parts(
-            self, transitions: Transitions) -> list[dict[str, float]]:
-        """Return [site occupancy][gemdat.sites.SitesData.site_occupancy] dicts
-        per part."""
-        labels = self.site_labels
-        n_steps = len(self.trajectory)
-
-        parts = transitions.split(self.n_parts)
-
-        return [
-            _calculate_site_occupancy(occupancy=part.occupancy(),
-                                      labels=labels,
-                                      n_steps=int(n_steps / self.n_parts))
-            for part in parts
-        ]
-
-    @weak_lru_cache()
-    def atom_locations_parts(
-            self, transitions: Transitions) -> list[dict[str, float]]:
-        """Return [atom locations][gemdat.sites.SitesData.atom_locations] dicts
-        per part."""
-        multiplier = self.n_sites / self.n_floating
-        return [{
-            k: v * multiplier
-            for k, v in part.items()
-        } for part in self.site_occupancy_parts(transitions)]
-
-    def site_occupancy(self, transitions: Transitions):
-        """Calculate percentage occupancy per unique site.
-
-        Returns
-        -------
-        site_occopancy : dict[str, float]
-            Percentage occupancy per unique site
-        """
-        labels = self.site_labels
-        n_steps = len(self.trajectory)
-        return _calculate_site_occupancy(occupancy=transitions.occupancy(),
-                                         labels=labels,
-                                         n_steps=n_steps)
-
-    def atom_locations(self, transitions):
+    def atom_locations(self, transitions: Transitions):
         """Calculate fraction of time atoms spent at a type of site.
 
         Returns
@@ -157,45 +103,15 @@ def atom_locations(self, transitions):
             Return dict with the fraction of time atoms spent at a site
         """
         multiplier = self.n_sites / self.n_floating
-        return {
-            k: v * multiplier
-            for k, v in self.site_occupancy(transitions).items()
-        }
-
-
-def _calculate_site_occupancy(
-    *,
-    occupancy: dict[int, int],
-    labels: list[str],
-    n_steps: int,
-) -> dict[str, float]:
-    """Calculate percentage occupancy per unique site.
-
-    Parameters
-    ----------
-    occupancy : dict[int, int]
-        Occupancy dict
-    labels : list[str]
-        Site labels
-    n_steps : int
-        Number of steps in time series
-
-    Returns
-    -------
-    dict[str, float]
-        Percentage occupancy per unique site
-    """
-    counts = defaultdict(list)
-
-    assert all(v >= 0 for v in occupancy)
-
-    for k, v in occupancy.items():
-        label = labels[k]
-        counts[label].append(v)
-
-    site_occupancies = {
-        k: sum(v) / (n_steps * labels.count(k))
-        for k, v in counts.items()
-    }
-
-    return site_occupancies
+
+        compositions_by_label = defaultdict(list)
+
+        for site in transitions.occupancy():
+            compositions_by_label[site.label].append(site.species.num_atoms)
+
+        ret = {}
+
+        for k, v in compositions_by_label.items():
+            ret[k] = (sum(v) / len(v)) * multiplier
+
+        return ret

src/gemdat/transitions.py

@@ -166,15 +166,32 @@ def states_prev(self) -> np.ndarray:
         """
         return ffill(self.states, fill_val=NOSITE, axis=0)
 
-    def occupancy(self) -> dict[int, int]:
+    def occupancy(self) -> Structure:
         """Calculate occupancy per site.
 
         Returns
         -------
-        occupancy : dict[int, int]
-            For each site, count for how many time steps it is occupied by an atom
+        structure : Structure
+            Structure with occupancies set on the sites.
         """
-        return _calculate_occupancy(self.states)
+        structure = self.structure
+        states = self.states
+
+        unq, counts = np.unique(states, return_counts=True)
+        counts = counts / len(states)
+        occupancies = dict(zip(unq, counts))
+
+        species = [{
+            site.specie.name: occupancies.get(i, 1)
+        } for i, site in enumerate(structure)]
+
+        return Structure(
+            lattice=structure.lattice,
+            species=species,
+            coords=structure.frac_coords,
+            site_properties=structure.site_properties,
+            labels=structure.labels,
+        )
 
     def split(self, n_parts: int = 10) -> list[Transitions]:
         """Split data into equal parts in time for statistics.
@@ -464,22 +481,3 @@ def _split_transitions_events(events: pd.DataFrame,
         part[dependent_keys] -= offset
 
     return parts
-
-
-def _calculate_occupancy(atom_sites: np.ndarray) -> dict[int, int]:
-    """Calculate occupancy per site.
-
-    Parameters
-    ----------
-    atom_sites : np.ndarray
-        Input array with atom sites
-
-    Returns
-    -------
-    occupancy : dict[int, int]
-        For each site, count for how many time steps it is occupied by an atom
-    """
-    unq, counts = np.unique(atom_sites, return_counts=True)
-    occupancy = dict(zip(unq, counts))
-    occupancy.pop(NOSITE, None)
-    return occupancy

tests/integration/sites_test.py

@@ -71,60 +71,24 @@ def test_transitions_parts(self, vasp_sites, vasp_transitions):
             np.array([[0, 0, 95], [4, 23, 95], [8, 28, 95]]))
 
     def test_occupancy(self, vasp_transitions):
-        occupancy = vasp_transitions.occupancy()
-        assert len(occupancy) == 95
-        assert sum(occupancy.values()) == 137026
-        assert list(occupancy.values())[::20] == [1704, 971, 351, 1508, 1104]
-
-    def test_occupancy_parts(self, vasp_sites, vasp_transitions):
-        parts = vasp_sites.occupancy_parts(vasp_transitions)
-        assert len(parts) == self.n_parts
-        assert [sum(part.values()) for part in parts] == [
-            13592, 13898, 13819, 14028, 14022, 14470, 13200, 13419, 13286,
-            13292
-        ]
-
-    def test_site_occupancy(self, vasp_sites, vasp_transitions):
-        assert isclose(vasp_sites.site_occupancy(vasp_transitions)['48h'],
-                       0.380628,
-                       rel_tol=1e-4)
+        structure = vasp_transitions.occupancy()
 
-    def test_site_occupancy_parts(self, vasp_sites, vasp_transitions):
-        assert len(
-            vasp_sites.site_occupancy_parts(vasp_transitions)) == self.n_parts
-        assert isclose(
-            vasp_sites.site_occupancy_parts(vasp_transitions)[0]['48h'],
-            0.37756,
-            rel_tol=1e-4)
-        assert isclose(
-            vasp_sites.site_occupancy_parts(vasp_transitions)[9]['48h'],
-            0.36922,
-            rel_tol=1e-4)
+        assert len(structure) == 96
+        assert structure[0].species.num_atoms == 0.4544
+        assert structure.composition.num_atoms == 37.54026666666666
 
     def test_atom_locations(self, vasp_sites, vasp_transitions):
         assert isclose(vasp_sites.atom_locations(vasp_transitions)['48h'],
-                       0.761255,
+                       0.7820889,
                        rel_tol=1e-4)
 
-    def test_atom_locations_parts(self, vasp_sites, vasp_transitions):
-        assert len(
-            vasp_sites.atom_locations_parts(vasp_transitions)) == self.n_parts
-        assert isclose(
-            vasp_sites.atom_locations_parts(vasp_transitions)[0]['48h'],
-            0.755111,
-            rel_tol=1e-4)
-        assert isclose(
-            vasp_sites.atom_locations_parts(vasp_transitions)[9]['48h'],
-            0.738444,
-            rel_tol=1e-4)
-
     def test_n_jumps(self, vasp_jumps):
         assert vasp_jumps.n_solo_jumps == 450
         assert vasp_jumps.n_jumps == 462
         assert isclose(vasp_jumps.solo_fraction, 0.974026, abs_tol=1e-4)
 
     def test_rates(self, vasp_jumps):
-        rates = vasp_jumps.rates()
+        rates = vasp_jumps.rates(n_parts=10)
         assert isinstance(rates, dict)
         assert len(rates) == 1
 
@@ -133,15 +97,15 @@ def test_rates(self, vasp_jumps):
         assert isclose(rates_std, 41893421993.683655)
 
     def test_activation_energies(self, vasp_jumps, vasp_sites):
-        activation_energies = vasp_jumps.activation_energies()
+        activation_energies = vasp_jumps.activation_energies(n_parts=10)
 
         assert isinstance(activation_energies, dict)
         assert len(activation_energies) == 1
 
         e_act, e_act_std = activation_energies[('48h', '48h')]
 
-        assert isclose(e_act, 0.1744591, abs_tol=1e-4)
-        assert isclose(e_act_std, .00405951, abs_tol=1e-6)
+        assert isclose(e_act, 0.20223, abs_tol=1e-4)
+        assert isclose(e_act_std, 0.00595, abs_tol=1e-6)
 
     def test_jump_diffusivity(self, vasp_jumps):
         assert isclose(vasp_jumps.jump_diffusivity(3),