Repository: GEMDAT-repos/GEMDAT (language: python)

Issue:
Refactor ported code
As we ported more, we realized that some structures could be better represented. this issue is to track that.

- [x] Refactor all calculate functions into "property" functions, which are calculated on demand
  - [x] Also add a toggle to calculate on creation of the object. 
- [x] move diffusive element displacement code into Trajectories
- [x] Remove calculate folder and put calculations in corresponding files
- [x] Refactor all plots to use the Trajectory / Sites / Jumps classes where possible, so that we do not have to pass each property independently 
- [x] Math ( or TrajectoryStatistics) class for statistics concerning Trajectories
- [x] remove diffusive_element specifiers on `Trajectory` functions, and instead allow a view on trajectory, like:
```python
li_trajectory = trajectory.where(element='Li', equilibration_steps=1250)
```
- [x] Better way to include all plots then using `__all__`
- [x] write `trajectory.precompute()` and `sites.precompute()`
- [x] Fix the RDF functions
Refactor GEMDAT
- [ ] Implement everything from issue #93 
- [x] Fix all the tests except rdf
- [ ]  Fix the RDF tests
- [x] Test the dashboard

Script I am working with:
```python
from pathlib import Path

from gemdat import Trajectory, plot, Vibration
from gemdat.sites import SitesData
from gemdat.io import load_known_material
import matplotlib.pyplot as plt

trajectory = Trajectory.from_vasprun(Path('../short_simulation/vasprun.xml'))
trajectory = trajectory[1250:]
li_trajectory = trajectory.filter('Li')
structure = load_known_material('argyrodite', supercell=(2,1,1))
li_vibration = Vibration(li_trajectory, fs=1.)

sd = SitesData(structure, li_trajectory, li_vibration)
plot(trajectory=li_trajectory,plots=["displacement_per_site"],show=False)
plot(trajectory=trajectory,plots=["displacement_per_element"],show=False)
plot(trajectory=li_trajectory,plots=["displacement_histogram"],show=False)
plot(plots=["frequency_vs_occurence"],vibration=li_vibration, show=False)
plot(plots=["vibrational_amplitudes"],vibration=li_vibration, show=False)
plot(plots=["jumps_vs_distance"],trajectory=li_trajectory, sites=sd, jump_res=1.0, show=False)
plot(plots=["jumps_vs_time"],trajectory=li_trajectory, sites=sd, show=False)
plot(plots=["collective_jumps"],trajectory=li_trajectory, sites=sd, show=False)
plot(plots=["jumps_3d"],trajectory=li_trajectory, sites=sd, show=False)
plt.show()
```


Discussion / hints:

@stefsmeets Curious what you think about this, 
The changes are quite big, but the numbers fortunately stay almost the same.

I removed the calculate all function and used (cached) class functions to calculate arrays. a `precompute` still needs to be implemented that calls each function once, but its not critical.

I especially like the added `parts` and `filter` function for the `Trajectory`. I am not yet sure it is a very good idea to have functional access to all arrays, although it seems to work nicely.
Hi @v1kko This PR seems like it just changes about everything in gemdat, and I'm not entirely sure whether I am onboard with all the changes. What exactly is the goal of this refactor?

This PR is also very large, which makes it difficult for me to review it. Just having a glance, I see lots of small issues. If the PR was smaller, it would be easier for me to go through it without losing focus.

Some comments:
- I looked at the Trajectory refactor, and I don't like where the design of the API is headed in this PR. There are a bunch of methods mixed in there that having nothing to do with a Trajectory in my opinion. Things like `tracer_diffusivity` and `mol_per_liter` make no sense there. I thought we wanted to split those out, but instead it feels like we are circling back to what we had in the first place.
- Is this PR 'finished'? It feels a bit rough around the edges to me looking at the docstrings / type hints / comments etc. There are too many things I want to comment on, it would be easier if this PR was split up. 
- `Vibration.meanfreq` would be better as a method. It has nothing to do with `Vibration` in particular, because it mimicks the matlab function: https://nl.mathworks.com/help/signal/ref/meanfreq.html.
- Some of the differences in values in the integration test are concerning. I would expect all test values to stay exactly the same.
- `Sites.parts` does not work this way, because it does not track the atom state (start / end site) correctly if it the split happens when it is between sites.
Hi @stefsmeets, 

- This PR is not finished, but I wanted to hear your thoughts on it.
- I agree that  the tracer/mol_per_liter part should be factored into its own math/utilities class, something that can be done quite quickly now.
- The difference for the integration test would be most likely because of the `Sites.parts` function, as most (all?) differences are in that area. I expected that a recomputation from the splitted trajectories would yield the same result, but as you say, it doesn't. The idea behind this was that with recomputation the arrays would not have to be split manually.

I can try to split this pull request in smaller ones, but that will take quite some time
Yeah, I fell into the same trap last week with `sites.parts`, where I naively thought I could just use the split trajectory directly. I was wrong 😅

That's why we need to derive everything from `all_transitions` / `atom_sites{_to,_from}` instead.

Fix patch:
diff --git a/README.md b/README.md
index 5b9d8f29..cdaabcf3 100644
--- a/README.md
+++ b/README.md
@@ -30,21 +30,29 @@ Suggestions, improvements, and edits are most welcome.
 The following snippet can be used to test the code using VASP data.
 
 ```python
-from gemdat import Trajectory, plot_all, plot
+from gemdat import Trajectory, plots
 from gemdat.calculate.extras import calculate_all
 
 trajectory = Trajectory.from_vasprun(Path('../example/vasprun.xml'))
-extras = calculate_all(diffusing_element='Li')
+
+diff_trajectory = trajectory.filter('Li')
+
+plots.plot_displacement_per_element(trajectory)
+plots.plot_displacement_per_site(diff_trajectory)
+plots.plot_displacement_histogram(diff_trajectory)
+plots.plot_frequency_vs_occurence(trajectory)
+plots.plot_vibrational_amplitudes(trajectory)
 
 structure = load_known_material('argyrodite', supercell=(2, 1, 1))
 
 sites = SitesData(structure)
 sites.calculate_all(trajectory=trajectory,
-                    extras=extras)
+                    diffusing_element='Li')
 
-plot_all(trajectory=trajectory,
-         sites=sites,
-         **vars(extras),
+plots.plot_jumps_vs_distance(trajectory, sites)
+plots.plot_jumps_vs_time(trajectory, sites)
+plots.plot_collective_jumps(trajectory, sites)
+plots.plot_jumps_3d(trajectory, sites)
 ```
 
 Or, one function to do everything:
diff --git a/src/calculate/__init__.py b/src/calculate/__init__.py
deleted file mode 100644
index cd33ef06..00000000
--- a/src/calculate/__init__.py
+++ /dev/null
@@ -1,10 +0,0 @@
-from .extras import calculate_all
-from .tracer import Tracer
-from .vibration import Vibration
-
-__all__ = [
-    'calculate_all',
-    'Displacements',
-    'Vibration',
-    'Tracer',
-]
diff --git a/src/calculate/extras.py b/src/calculate/extras.py
deleted file mode 100644
index 2a17d4d3..00000000
--- a/src/calculate/extras.py
+++ /dev/null
@@ -1,59 +0,0 @@
-from __future__ import annotations
-
-import typing
-from types import SimpleNamespace
-
-from .tracer import Tracer
-from .vibration import Vibration
-
-if typing.TYPE_CHECKING:
-    from gemdat.trajectory import Trajectory
-
-
-def calculate_all(trajectory: Trajectory,
-                  *,
-                  diffusing_element: str,
-                  known_structure: str | None = None,
-                  diffusion_dimensions: int = 3,
-                  z_ion: float = 1.0,
-                  n_parts: int = 10,
-                  dist_collective: float = 4.5):
-    """Calculate extra parameters and return them as a simple namespace.
-
-    Parameters
-    ----------
-    trajectory : Trajectory
-        Input trajectory coordinates
-    diffusing_element : str
-        Name of the diffusing element
-    structure : str | None
-        Path to cif file or name of known structure
-    diffusion_dimensions : int
-        Number of diffusion dimensions
-    z_ion : float
-        Ionic charge of the diffusing ion
-    n_parts : int
-        In how many parts to divide your simulation for statistics
-    dist_collective : float
-        Maximum distance for collective motions in Angstrom
-    """
-    _ = trajectory.to_displacements()
-
-    diff_trajectory = trajectory.filter(diffusing_element)
-
-    extras = SimpleNamespace(
-        diffusing_element=diffusing_element,
-        known_structure=known_structure,
-        diffusion_dimensions=diffusion_dimensions,
-        z_ion=z_ion,
-        n_parts=n_parts,
-        dist_collective=dist_collective,
-        n_steps=len(trajectory),
-        total_time=len(trajectory) * trajectory.time_step,
-    )
-
-    extras.__dict__.update(Vibration.calculate_all(diff_trajectory))
-    extras.__dict__.update(Tracer.calculate_all(diff_trajectory,
-                                                extras=extras))
-
-    return extras
diff --git a/src/calculate/tracer.py b/src/calculate/tracer.py
deleted file mode 100644
index 1f99cdcf..00000000
--- a/src/calculate/tracer.py
+++ /dev/null
@@ -1,72 +0,0 @@
-from __future__ import annotations
-
-import typing
-
-import numpy as np
-from pymatgen.core.units import FloatWithUnit
-from scipy.constants import Avogadro, Boltzmann, angstrom, elementary_charge
-
-if typing.TYPE_CHECKING:
-    from types import SimpleNamespace
-
-    from gemdat.trajectory import Trajectory
-
-
-class Tracer:
-
-    @staticmethod
-    def calculate_all(trajectory: Trajectory,
-                      extras: SimpleNamespace) -> dict[str, float]:
-        """Calculate tracer properties.
-
-        Parameters
-        ----------
-        trajectory : Trajectory
-            Input trajectory for diffusing species
-        extras : SimpleNamespace
-            Extra variables
-
-        Returns
-        -------
-        extras : dict[str, float]
-            Dictionary with calculated parameters
-        """
-        lattice = trajectory.get_lattice()
-
-        total_time = extras.total_time
-
-        volume_ang = lattice.volume
-        volume_m3 = volume_ang * angstrom**3
-
-        particle_density = len(trajectory.species) / volume_m3
-
-        mol_per_liter = (particle_density * 1e-3) / Avogadro
-
-        particle_density = FloatWithUnit(particle_density, 'm^-3')
-        mol_per_liter = FloatWithUnit(mol_per_liter, 'mol l^-1')
-
-        # Matlab code contains a bug here so I'm not entirely sure what is the definition
-        # Matlab code takes the first column, which is equal to 0
-        # Do they mean the total displacement (i.e. last column)?
-        msd = np.mean(
-            trajectory.distances_from_base_position()[:, -1]**2)  # Angstrom^2
-
-        temperature = trajectory.metadata['temperature']
-
-        # Diffusivity = MSD/(2*dimensions*time)
-        tracer_diff = (msd * angstrom**2) / (2 * extras.diffusion_dimensions *
-                                             total_time)
-        # Conductivity = elementary_charge^2 * charge_ion^2 * diffusivity * particle_density / (k_B * T)
-        tracer_conduc = ((elementary_charge**2) *
-                         (extras.z_ion**2) * tracer_diff *
-                         particle_density) / (Boltzmann * temperature)
-
-        tracer_diff = FloatWithUnit(tracer_diff, 'm^2 s^-1')
-        tracer_conduc = FloatWithUnit(tracer_conduc, 'S m^-1')
-
-        return {
-            'particle_density': particle_density,
-            'mol_per_liter': mol_per_liter,
-            'tracer_diff': tracer_diff,
-            'tracer_conduc': tracer_conduc,
-        }
diff --git a/src/calculate/vibration.py b/src/calculate/vibration.py
deleted file mode 100644
index 1672547a..00000000
--- a/src/calculate/vibration.py
+++ /dev/null
@@ -1,178 +0,0 @@
-from __future__ import annotations
-
-import typing
-
-import numpy as np
-from pymatgen.core.units import FloatWithUnit
-from scipy import signal
-
-if typing.TYPE_CHECKING:
-
-    from gemdat.trajectory import Trajectory
-
-
-def meanfreq(x: np.ndarray, fs: float = 1.0):
-    """Estimates the mean frequency in terms of the sample rate, fs.
-
-    Vectorized version of https://stackoverflow.com/a/56487241
-
-    Parameters
-    ----------
-    x : np.ndarray[i, j]
-        Time series of measurement values. The mean frequency is computed
-        along the last axis (-1).
-    fs : float, optional
-        Sampling frequency of the `x` time series. Defaults to 1.0.
-
-    Returns
-    -------
-    mnfreq : np.ndarray
-        Array of mean frequencies.
-    """
-    if x.ndim == 1:
-        x = x.reshape(1, -1)
-
-    assert x.ndim == 2
-
-    f, Pxx_den = signal.periodogram(x, fs, axis=-1)
-    width = np.tile(f[1] - f[0], Pxx_den.shape)
-    P = Pxx_den * width
-    pwr = np.sum(P, axis=1).reshape(-1, 1)
-
-    f = f.reshape(1, -1)
-
-    mnfreq = np.dot(P, f.T) / pwr
-
-    return mnfreq
-
-
-class Vibration:
-
-    @staticmethod
-    def calculate_all(trajectory: Trajectory) -> dict:
-        """Calculate Vibration properties.
-
-        Parameters
-        ----------
-        trajectory : Trajectory
-            Input trajectory for diffusing species
-
-        Returns
-        -------
-        extras : dict[str, float]
-            Dictionary with calculated parameters
-        """
-        fs = 1 / trajectory.time_step
-
-        speed = Vibration.speed(trajectory.distances_from_base_position())
-        attempt_freq, attempt_freq_std = Vibration.attempt_frequency(speed,
-                                                                     fs=fs)
-
-        amplitudes = Vibration.amplitudes(speed)
-        vibration_amplitude = Vibration.vibration_amplitude(amplitudes)
-
-        return {
-            'speed': speed,
-            'attempt_freq': attempt_freq,
-            'attempt_freq_std': attempt_freq_std,
-            'amplitudes': amplitudes,
-            'vibration_amplitude': vibration_amplitude,
-            'fs': fs,
-        }
-
-    @staticmethod
-    def speed(distances: np.ndarray) -> np.ndarray:
-        """Calculate attempt frequency.
-
-        Parameters
-        ----------
-        distances : np.ndarray
-            Input array with distances
-
-        Returns
-        -------
-        speed : np.ndarray
-            Output array with speeds
-        """
-        return np.diff(distances, prepend=0)
-
-    @staticmethod
-    def attempt_frequency(speed: np.ndarray,
-                          fs: float = 1) -> tuple[float, float]:
-        """Calculate attempt frequency.
-
-        Parameters
-        ----------
-        speed : np.ndarray
-            Input array with displacement speeds
-        fs : float, optional
-            Sampling frequency
-
-        Returns
-        -------
-        attempt_freq : float
-            Attempt frequency
-        attempt_freq_std : float
-            Attempt frequency standard deviation
-        """
-        freq_mean = meanfreq(speed, fs=fs)
-
-        attempt_freq = np.mean(freq_mean)
-        attempt_freq_std = np.std(freq_mean)
-
-        attempt_freq = FloatWithUnit(attempt_freq, 'hz')
-        attempt_freq_std = FloatWithUnit(attempt_freq_std, 'hz')
-
-        return attempt_freq, attempt_freq_std
-
-    @staticmethod
-    def amplitudes(speed: np.ndarray) -> np.ndarray:
-        """Calculate vibration amplitude.
-
-        Parameters
-        ----------
-        speed : np.ndarray
-            Input array with displacement velocities
-
-        Returns
-        -------
-        amplitudes : np.ndarray
-            Output array with vibrational amplitudes
-        vibration amplitude : float
-            Mean vibration amplitude
-        """
-        amplitudes = []
-
-        for i, speed_range in enumerate(speed):
-            signs = np.sign(speed_range)
-
-            # get indices where sign flips
-            splits = np.where(signs != np.roll(signs, shift=-1))[0]
-            # strip first and last splits
-            subarrays = np.split(speed_range, splits[1:-1] + 1)
-
-            amplitudes.extend([np.sum(array) for array in subarrays])
-
-        return np.asarray(amplitudes)
-
-    @staticmethod
-    def vibration_amplitude(amplitudes: np.ndarray) -> float:
-        """Calcualte vibration ambplitude.
-
-        Parameters
-        ----------
-        amplitudes : np.ndarray
-            Input amplitudes
-
-        Returns
-        -------
-        vibration_amp : float
-            Vibration amplitude
-        """
-        mean_vib = np.mean(amplitudes)
-        vibration_amp: float = np.std(amplitudes)
-
-        mean_vib = FloatWithUnit(mean_vib, 'ang')
-        vibration_amp = FloatWithUnit(vibration_amp, 'ang')
-
-        return vibration_amp
diff --git a/src/dashboard/run.py b/src/dashboard/run.py
index 945ea420..b2c3b458 100644
--- a/src/dashboard/run.py
+++ b/src/dashboard/run.py
@@ -1,9 +1,9 @@
 import streamlit as st
 from _shared import add_sidebar_logo, get_trajectory_location
 from gemdat import SitesData, __version__, plots
-from gemdat.calculate import calculate_all
 from gemdat.io import get_list_of_known_materials, load_known_material
 from gemdat.rdf import calculate_rdfs, plot_rdf
+from gemdat.simulation_metrics import SimulationMetrics
 from gemdat.trajectory import Trajectory
 from gemdat.utils import is_lattice_similar
 
@@ -78,29 +78,31 @@ def _load_trajectory(trajectory_location):
 
 number_of_cols = 3  # Number of figure columns
 
-with st.spinner('Processing trajectory...'):
-    trajectory = trajectory[equilibration_steps:]
-    extras = calculate_all(trajectory, diffusing_element=diffusing_element)
+trajectory = trajectory[equilibration_steps:]
+diff_trajectory = trajectory.filter('Li')
+metrics = SimulationMetrics(diff_trajectory)
 
 col1, col2, col3 = st.columns(3)
 
 with col1:
     import uncertainties as u
-    attempt_freq = u.ufloat(extras.attempt_freq, extras.attempt_freq_std)
+    attempt_freq = u.ufloat(metrics.attempt_frequency())
     st.metric('Attempt frequency ($\\mathrm{s^{-1}}$)',
               value=f'{attempt_freq:g}')
     st.metric('Vibration amplitude ($\\mathrm{Å}$)',
-              value=f'{extras.vibration_amplitude:g}')
+              value=f'{metrics.vibration_amplitude():g}')
 with col2:
     st.metric('Particle density ($\\mathrm{m^{-3}}$)',
-              value=f'{extras.particle_density:g}')
+              value=f'{metrics.particle_density():g}')
     st.metric('Mol per liter ($\\mathrm{mol/l}$)',
-              value=f'{extras.mol_per_liter:g}')
+              value=f'{metrics.mol_per_liter():g}')
 with col3:
     st.metric('Tracer diffusivity ($\\mathrm{m^2/s}$)',
-              value=f'{extras.tracer_diff:g}')
-    st.metric('Tracer conductivity ($\\mathrm{S/m}$)',
-              value=f'{extras.tracer_conduc:g}')
+              value=f'{metrics.tracer_diffusivity(diffusion_dimensions=3):g}')
+    st.metric(
+        'Tracer conductivity ($\\mathrm{S/m}$)',
+        value=
+        f'{metrics.tracer_conductivity(z_ion=1, diffusion_dimensions=3):g}')
 
 tab1, tab2 = st.tabs(['Default plots', 'RDF plots'])
 
@@ -119,7 +121,13 @@ def _load_trajectory(trajectory_location):
 
     with st.spinner('Calculating jumps...'):
         sites = SitesData(sites_structure)
-        sites.calculate_all(trajectory=trajectory, extras=extras)
+        sites.calculate_all(
+            trajectory=trajectory,
+            diffusing_element=diffusing_element,
+            z_ion=1,
+            diffusion_dimensions=3,
+            n_parts=10,
+        )
 
     diff_trajectory = trajectory.filter(diffusing_element)
 
@@ -127,12 +135,8 @@ def _load_trajectory(trajectory_location):
         plots.plot_displacement_per_element(trajectory=trajectory),
         plots.plot_displacement_per_site(trajectory=diff_trajectory),
         plots.plot_displacement_histogram(trajectory=diff_trajectory),
-        plots.plot_frequency_vs_occurence(trajectory=trajectory,
-                                          sites=sites,
-                                          **vars(extras)),
-        plots.plot_vibrational_amplitudes(trajectory=trajectory,
-                                          sites=sites,
-                                          **vars(extras)),
+        plots.plot_frequency_vs_occurence(trajectory=trajectory),
+        plots.plot_vibrational_amplitudes(trajectory=trajectory),
         plots.plot_jumps_vs_distance(trajectory=trajectory, sites=sites),
         plots.plot_jumps_vs_time(trajectory=trajectory, sites=sites),
         plots.plot_collective_jumps(trajectory=trajectory, sites=sites),
diff --git a/src/legacy.py b/src/legacy.py
index e3d97686..4dda7101 100644
--- a/src/legacy.py
+++ b/src/legacy.py
@@ -1,7 +1,4 @@
-from types import SimpleNamespace
-
 from gemdat import SitesData, load_known_material, plots
-from gemdat.calculate import calculate_all
 from gemdat.rdf import calculate_rdfs, plot_rdf
 from gemdat.trajectory import Trajectory
 from gemdat.volume import trajectory_to_vasp_volume
@@ -25,7 +22,7 @@ def analyse_md(
     rdf_max_dist: int = 10,
     start_end: tuple[int, int] = (5000, 7500),
     nr_steps_frame: int = 5,
-) -> tuple[Trajectory, SitesData, SimpleNamespace]:
+) -> tuple[Trajectory, SitesData]:
     """Analyse md data.
 
     Parameters
@@ -76,30 +73,23 @@ def analyse_md(
 
     diff_trajectory = trajectory.filter(diff_elem)
 
-    extras = calculate_all(
-        trajectory,
+    sites_structure = load_known_material(material, supercell=supercell)
+
+    sites = SitesData(sites_structure)
+    sites.calculate_all(
+        trajectory=trajectory,
         diffusing_element=diff_elem,
-        z_ion=z_ion,
         diffusion_dimensions=diffusion_dimensions,
+        z_ion=z_ion,
         n_parts=nr_parts,
+        dist_collective=dist_collective,
     )
 
-    sites_structure = load_known_material(material, supercell=supercell)
-
-    sites = SitesData(sites_structure)
-    sites.calculate_all(trajectory=trajectory,
-                        extras=extras,
-                        dist_collective=dist_collective)
-
     plots.plot_displacement_per_element(trajectory=trajectory)
     plots.plot_displacement_per_site(trajectory=diff_trajectory)
     plots.plot_displacement_histogram(trajectory=diff_trajectory)
-    plots.plot_frequency_vs_occurence(trajectory=trajectory,
-                                      sites=sites,
-                                      **vars(extras))
-    plots.plot_vibrational_amplitudes(trajectory=trajectory,
-                                      sites=sites,
-                                      **vars(extras))
+    plots.plot_frequency_vs_occurence(trajectory=trajectory)
+    plots.plot_vibrational_amplitudes(trajectory=trajectory)
     plots.plot_jumps_vs_distance(trajectory=trajectory,
                                  sites=sites,
                                  jump_res=jump_res)
@@ -136,4 +126,4 @@ def analyse_md(
         for name, rdf in rdf_data.items():
             plot_rdf(rdf, name=name)
 
-    return trajectory, sites, extras
+    return trajectory, sites
diff --git a/src/plots/__init__.py b/src/plots/__init__.py
index aa4a3316..7a9eba99 100644
--- a/src/plots/__init__.py
+++ b/src/plots/__init__.py
@@ -3,7 +3,13 @@
     plot_displacement_per_element,
     plot_displacement_per_site,
 )
-from .jumps import plot_collective_jumps, plot_jumps_3d, plot_jumps_vs_distance, plot_jumps_vs_time
+from .jumps import (
+    plot_collective_jumps,
+    plot_jumps_3d,
+    plot_jumps_3d_animation,
+    plot_jumps_vs_distance,
+    plot_jumps_vs_time,
+)
 from .vibration import (
     plot_frequency_vs_occurence,
     plot_vibrational_amplitudes,
@@ -19,4 +25,5 @@
     'plot_jumps_vs_time',
     'plot_collective_jumps',
     'plot_jumps_3d',
+    'plot_jumps_3d_animation',
 ]
diff --git a/src/plots/displacements.py b/src/plots/displacements.py
index 8a525d13..af13b96e 100644
--- a/src/plots/displacements.py
+++ b/src/plots/displacements.py
@@ -3,16 +3,13 @@
 from gemdat.trajectory import Trajectory
 
 
-def plot_displacement_per_site(*, trajectory: Trajectory,
-                               **kwargs) -> plt.Figure:
+def plot_displacement_per_site(*, trajectory: Trajectory) -> plt.Figure:
     """Plot displacement per site.
 
     Parameters
     ----------
     trajectory : Trajectory
         Input trajectory
-    **kwargs : dict
-        Extra parameters
 
     Returns
     -------
@@ -31,16 +28,13 @@ def plot_displacement_per_site(*, trajectory: Trajectory,
     return fig
 
 
-def plot_displacement_per_element(*, trajectory: Trajectory,
-                                  **kwargs) -> plt.Figure:
+def plot_displacement_per_element(*, trajectory: Trajectory) -> plt.Figure:
     """Plot displacement per element.
 
     Parameters
     ----------
     trajectory : Trajectory
         Input trajectory
-    **kwargs : dict
-        Extra parameters
 
     Returns
     -------
@@ -71,8 +65,7 @@ def plot_displacement_per_element(*, trajectory: Trajectory,
     return fig
 
 
-def plot_displacement_histogram(trajectory: Trajectory,
-                                **kwargs) -> plt.Figure:
+def plot_displacement_histogram(trajectory: Trajectory) -> plt.Figure:
     """Plot histogram of total displacement of diffusing element at final
     timestep.
 
@@ -80,8 +73,6 @@ def plot_displacement_histogram(trajectory: Trajectory,
     ----------
     trajectory : Trajectory
         Input trajectory
-    **kwargs : dict
-        Extra parameters
 
     Returns
     -------
diff --git a/src/plots/jumps.py b/src/plots/jumps.py
index cbf5bc0e..acc6cfb1 100644
--- a/src/plots/jumps.py
+++ b/src/plots/jumps.py
@@ -16,8 +16,7 @@
 def plot_jumps_vs_distance(*,
                            trajectory: Trajectory,
                            sites: SitesData,
-                           jump_res: float = 0.1,
-                           **kwargs) -> plt.Figure:
+                           jump_res: float = 0.1) -> plt.Figure:
     """Plot jumps vs. distance histogram.
 
     Parameters
@@ -61,8 +60,7 @@ def plot_jumps_vs_distance(*,
 def plot_jumps_vs_time(*,
                        trajectory: Trajectory,
                        sites: SitesData,
-                       binsize: int = 500,
-                       **kwargs) -> plt.Figure:
+                       binsize: int = 500) -> plt.Figure:
     """Plot jumps vs. distance histogram.
 
     Parameters
@@ -74,7 +72,6 @@ def plot_jumps_vs_time(*,
     binsize : int, optional
         Width of each bin in number of time steps
 
-
     Returns
     -------
     plt.Figure
@@ -93,8 +90,8 @@ def plot_jumps_vs_time(*,
     return fig
 
 
-def plot_collective_jumps(*, trajectory: Trajectory, sites: SitesData,
-                          **kwargs) -> plt.Figure:
+def plot_collective_jumps(*, trajectory: Trajectory,
+                          sites: SitesData) -> plt.Figure:
     """Plot collective jumps per jump-type combination.
 
     Parameters
@@ -124,8 +121,7 @@ def plot_collective_jumps(*, trajectory: Trajectory, sites: SitesData,
     return fig
 
 
-def plot_jumps_3d(*, trajectory: Trajectory, sites: SitesData,
-                  **kwargs) -> plt.Figure:
+def plot_jumps_3d(*, trajectory: Trajectory, sites: SitesData) -> plt.Figure:
     """Plot jumps in 3D.
 
     Parameters
@@ -214,8 +210,7 @@ def plot_jumps_3d_animation(*,
                             t_stop: int,
                             decay: float = 0.05,
                             skip: int = 5,
-                            interval: int = 20,
-                            **kwargs):
+                            interval: int = 20):
     """Plot jumps in 3D as an animation over time.
 
     # TODO
diff --git a/src/plots/vibration.py b/src/plots/vibration.py
index 3f473d06..17e5b5a9 100644
--- a/src/plots/vibration.py
+++ b/src/plots/vibration.py
@@ -1,29 +1,25 @@
-from typing import Optional
-
 import matplotlib.pyplot as plt
 import numpy as np
+from gemdat.simulation_metrics import SimulationMetrics
+from gemdat.trajectory import Trajectory
 from scipy import stats
 
 
-def plot_frequency_vs_occurence(*,
-                                speed: np.ndarray,
-                                fs: float,
-                                attempt_freq: Optional[float] = None,
-                                attempt_freq_std: Optional[float] = None,
-                                **kwargs):
+def plot_frequency_vs_occurence(*, trajectory: Trajectory) -> plt.Figure:
     """Plot attempt frequency vs occurence.
 
     Parameters
     ----------
-    speed : np.ndarray
-        Input array with displacement velocities
-    fs : float, optional
-        Sampling frequency
-    freq : Optional[float], optional
-        Attempt frequency
-    freq_std : Optional[float], optional
-        Attempt frequency standard deviation
+    trajectory : Trajectory
+        Input trajectory
+
+    Returns
+    -------
+    fig : plt.Figure
+        Output figure
     """
+    metrics = SimulationMetrics(trajectory)
+    speed = metrics.speed()
 
     length = speed.shape[1]
     half_length = length // 2 + 1
@@ -35,7 +31,7 @@ def plot_frequency_vs_occurence(*,
 
     fig, ax = plt.subplots()
 
-    f = fs * np.arange(half_length) / length
+    f = trajectory.sampling_frequency * np.arange(half_length) / length
 
     sum_freqs = np.sum(one_sided, axis=0)
     smoothed = np.convolve(sum_freqs, np.ones(51), 'same') / 51
@@ -43,6 +39,8 @@ def plot_frequency_vs_occurence(*,
 
     y_max = np.max(sum_freqs)
 
+    attempt_freq, attempt_freq_std = metrics.attempt_frequency()
+
     if attempt_freq:
         ax.vlines([attempt_freq], 0, y_max, colors='red')
     if attempt_freq and attempt_freq_std:
@@ -63,23 +61,26 @@ def plot_frequency_vs_occurence(*,
     return fig
 
 
-def plot_vibrational_amplitudes(*, amplitudes: np.ndarray,
-                                vibration_amplitude: float, **kwargs):
+def plot_vibrational_amplitudes(*, trajectory: Trajectory) -> plt.Figure:
     """Plot histogram of vibrational amplitudes with fitted Gaussian.
 
     Parameters
     ----------
-    amplitudes : np.ndarray
-        Input array with vibrational amplitudes
-    vibration_amplitude : float
-        Sigma of the vibrational amplitudes
+    trajectory : Trajectory
+        Input trajectory
+
+    Returns
+    -------
+    fig : plt.Figure
+        Output figure
     """
+    metrics = SimulationMetrics(trajectory)
 
     fig, ax = plt.subplots()
-    ax.hist(amplitudes, bins=100, density=True)
+    ax.hist(metrics.amplitudes(), bins=100, density=True)
 
     x = np.linspace(-2, 2, 100)
-    y_gauss = stats.norm.pdf(x, 0, vibration_amplitude)
+    y_gauss = stats.norm.pdf(x, 0, metrics.vibration_amplitude())
     ax.plot(x, y_gauss, 'r')
 
     ax.set(title='Histogram of vibrational amplitudes with fitted Gaussian',
diff --git a/src/simulation_metrics.py b/src/simulation_metrics.py
new file mode 100644
index 00000000..37615c0c
--- /dev/null
+++ b/src/simulation_metrics.py
@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import typing
+from functools import lru_cache
+
+import numpy as np
+from pymatgen.core.units import FloatWithUnit
+from scipy.constants import Avogadro, Boltzmann, angstrom, elementary_charge
+
+from .utils import meanfreq
+
+if typing.TYPE_CHECKING:
+    from gemdat.trajectory import Trajectory
+
+
+class SimulationMetrics:
+
+    def __init__(self, trajectory: Trajectory):
+        self.trajectory = trajectory
+
+    @lru_cache
+    def speed(self) -> np.ndarray:
+        """Calculate speed.
+
+        Corresponds to change in distance from the base position.
+        """
+        distances = self.trajectory.distances_from_base_position()
+        return np.diff(distances, prepend=0)
+
+    @lru_cache
+    def particle_density(self) -> float:
+        """Return number of particles per m3."""
+        lattice = self.trajectory.get_lattice()
+        volume_ang = lattice.volume
+        volume_m3 = volume_ang * angstrom**3
+        particle_density = len(self.trajectory.species) / volume_m3
+        return FloatWithUnit(particle_density, 'm^-3')
+
+    @lru_cache
+    def mol_per_liter(self) -> float:
+        """Return particle density as mol/liter."""
+        mol_per_liter = (self.particle_density() * 1e-3) / Avogadro
+        return FloatWithUnit(mol_per_liter, 'mol l^-1')
+
+    @lru_cache
+    def tracer_diffusivity(self, *, diffusion_dimensions: int) -> float:
+        """Return tracer diffusivity in m2/s.
+
+        Defined as: MSD/(2*dimensions*time)
+
+        Parameters
+        ----------
+        diffusion_dimensions : int
+            Number of diffusion dimensions
+
+        Returns
+        -------
+        tracer_diffusivity : float
+        """
+        distances = self.trajectory.distances_from_base_position()
+        msd = np.mean(distances[:, -1]**2)  # Angstrom^2
+
+        tracer_diff = (msd * angstrom**2) / (2 * diffusion_dimensions *
+                                             self.trajectory.total_time)
+
+        return FloatWithUnit(tracer_diff, 'm^2 s^-1')
+
+    @lru_cache
+    def tracer_conductivity(self, *, z_ion: int,
+                            diffusion_dimensions: int) -> float:
+        """Return tracer conductivity as S/m.
+
+        Defined as: elementary_charge^2 * charge_ion^2 * diffusivity *
+            particle_density / (k_B * T)
+
+        Parameters
+        ----------
+        z_ion : int
+            Charge of the ion
+        diffusion_dimensions : int
+            Number of diffusion dimensions
+
+        Returns
+        -------
+        tracer_conductivity : float
+        """
+        temperature = self.trajectory.metadata['temperature']
+        tracer_diff = self.tracer_diffusivity(
+            diffusion_dimensions=diffusion_dimensions)
+        tracer_conduc = ((elementary_charge**2) * (z_ion**2) * tracer_diff *
+                         self.particle_density()) / (Boltzmann * temperature)
+
+        return FloatWithUnit(tracer_conduc, 'S m^-1')
+
+    @lru_cache
+    def attempt_frequency(self) -> tuple[float, float]:
+        """Return attempt frequency and standard deviation in Hz.
+
+        Returns
+        -------
+        attempt_freq : float
+            Attempt frequency
+        attempt_freq_std : float
+            Attempt frequency standard deviation
+        """
+        speed = self.speed()
+
+        freq_mean = meanfreq(speed, fs=self.trajectory.sampling_frequency)
+
+        attempt_freq_std = np.std(freq_mean)
+        attempt_freq_std = FloatWithUnit(attempt_freq_std, 'hz')
+
+        attempt_freq = np.mean(freq_mean)
+        attempt_freq = FloatWithUnit(attempt_freq, 'hz')
+
+        return attempt_freq, attempt_freq_std
+
+    @lru_cache
+    def vibration_amplitude(self) -> float:
+        """Return vibration amplitude in Angstrom.
+
+        Returns
+        -------
+        vibration_amp : float
+        """
+        amplitudes = self.amplitudes()
+
+        mean_vib = np.mean(amplitudes)
+        vibration_amp = np.std(amplitudes)
+
+        mean_vib = FloatWithUnit(mean_vib, 'ang')
+        vibration_amp = FloatWithUnit(vibration_amp, 'ang')
+
+        return vibration_amp
+
+    @lru_cache
+    def amplitudes(self) -> np.ndarray:
+        """Return vibration amplitudes.
+
+        Returns
+        -------
+        amplitudes : np.ndarray
+            Output array of vibration amplitudes
+        """
+        amplitudes = []
+        speed = self.speed()
+
+        for i, speed_range in enumerate(speed):
+            signs = np.sign(speed_range)
+
+            # get indices where sign flips
+            splits = np.where(signs != np.roll(signs, shift=-1))[0]
+            # strip first and last splits
+            subarrays = np.split(speed_range, splits[1:-1] + 1)
+
+            amplitudes.extend([np.sum(array) for array in subarrays])
+
+        return np.asarray(amplitudes)
diff --git a/src/sites.py b/src/sites.py
index 02d5b3bd..8142bd7b 100644
--- a/src/sites.py
+++ b/src/sites.py
@@ -10,10 +10,10 @@
 from pymatgen.core.units import FloatWithUnit
 from scipy.constants import Boltzmann, angstrom, elementary_charge
 
+from .simulation_metrics import SimulationMetrics
 from .utils import bfill, ffill, is_lattice_similar
 
 if typing.TYPE_CHECKING:
-    from types import SimpleNamespace
 
     from gemdat.trajectory import Trajectory
 
@@ -46,26 +46,43 @@ def warn_if_lattice_not_similar(self, other_lattice: Lattice):
             warnings.warn(f'Lattice mismatch: {this_lattice.parameters} '
                           f'vs. {other_lattice.parameters}')
 
-    def calculate_all(self, trajectory: Trajectory, extras: SimpleNamespace,
-                      **kwargs):
+    def calculate_all(
+        self,
+        trajectory: Trajectory,
+        *,
+        diffusing_element: str,
+        z_ion: int = 1,
+        diffusion_dimensions: int = 3,
+        dist_collective: float = 4.5,
+        n_parts: int = 1,
+    ):
         """Calculate all parameters.
 
         Parameters
         ----------
         trajectory : Trajectory
             Input trajectory
-        extras : SimpleNamespace
-            Extra parameters
+        diffusing_element : str
+            Name of the diffusing species
+        z_ion : int
+            Ionic charge of difussing species
+        diffusion_dimensions : int
+            Number of diffusion dimensions
+        n_parts : int
+            Number of parts to divide simulation into for statistics
         """
         lattice = trajectory.get_lattice()
 
         self.warn_if_lattice_not_similar(lattice)
 
-        self.dist_close = self.calculate_dist_close(
-            trajectory, vibration_amplitude=extras.vibration_amplitude)
-
-        diff_trajectory = trajectory.filter(extras.diffusing_element)
+        diff_trajectory = trajectory.filter(diffusing_element)
         n_diffusing = len(diff_trajectory.species)
+        n_steps = len(diff_trajectory)
+
+        metrics = SimulationMetrics(diff_trajectory)
+
+        self.dist_close = self.calculate_dist_close(
+            trajectory, vibration_amplitude=metrics.vibration_amplitude())
 
         self.atom_sites = self.calculate_atom_sites(diff_trajectory)
         self.atom_sites_from = self.calculate_atom_sites_from()
@@ -77,8 +94,7 @@ def calculate_all(self, trajectory: Trajectory, extras: SimpleNamespace,
 
         self.occupancy = self.calculate_occupancy()
 
-        self.sites_occupancy = self.calculate_site_occupancy(
-            n_steps=extras.n_steps)
+        self.sites_occupancy = self.calculate_site_occupancy(n_steps=n_steps)
 
         self.atom_locations = self.calculate_atom_locations(
             n_diffusing=n_diffusing)
@@ -89,15 +105,20 @@ def calculate_all(self, trajectory: Trajectory, extras: SimpleNamespace,
         self.jump_diffusivity = self.calculate_jump_diffusivity(
             lattice=lattice,
             n_diffusing=n_diffusing,
-            total_time=extras.total_time,
-            dimensions=extras.diffusion_dimensions)
-        self.correlation_factor = extras.tracer_diff / self.jump_diffusivity
+            total_time=trajectory.total_time,
+            dimensions=diffusion_dimensions)
+        tracer_diff = metrics.tracer_diffusivity(
+            diffusion_dimensions=diffusion_dimensions)
+        self.correlation_factor = tracer_diff / self.jump_diffusivity
+
+        attempt_freq, _ = metrics.attempt_frequency()
 
         self.collective, self.coll_jumps, self.n_solo_jumps = self.calculate_collective(
             lattice=lattice,
-            attempt_freq=extras.attempt_freq,
+            attempt_freq=attempt_freq,
             time_step=trajectory.time_step,
-            **kwargs)
+            dist_collective=dist_collective,
+        )
 
         self.solo_frac = self.n_solo_jumps / len(self.all_transitions)
         self.coll_count = len(self.collective)
@@ -106,18 +127,20 @@ def calculate_all(self, trajectory: Trajectory, extras: SimpleNamespace,
         self.multi_coll = self.calculate_multiple_collective()
 
         # calculate some statistics
-        if extras.n_parts > 1:
-            self.parts = self.split(n_parts=extras.n_parts,
-                                    n_diffusing=n_diffusing,
-                                    extras=extras)
+        if n_parts > 1:
+            self.parts = self.split(
+                n_parts=n_parts,
+                n_steps=n_steps,
+                n_diffusing=n_diffusing,
+            )
 
-            self.rates = self.calculate_rates(total_time=extras.total_time,
+            self.rates = self.calculate_rates(total_time=trajectory.total_time,
                                               n_diffusing=n_diffusing)
 
             self.activation_energies = self.calculate_activation_energies(
-                total_time=extras.total_time,
+                total_time=trajectory.total_time,
                 n_diffusing=n_diffusing,
-                attempt_freq=extras.attempt_freq,
+                attempt_freq=attempt_freq,
                 temperature=trajectory.metadata['temperature'])
 
     @property
@@ -263,18 +286,18 @@ def calculate_occupancy(self):
         """
         return _calculate_occupancy(self.atom_sites)
 
-    def split(self, n_parts: int, n_diffusing: int,
-              extras: SimpleNamespace) -> list[SitesData]:
+    def split(self, *, n_parts: int, n_steps: int,
+              n_diffusing: int) -> list[SitesData]:
         """Split data into equal parts in time for internal statistics.
 
         Parameters
         ----------
         n_parts : int
             Number of parts to split the data into
+        n_steps : int
+            Number of steps in the trajectory
         n_diffusing : int
             Number of diffusing species
-        extras : SimpleNamespace
-            Extra parameters
 
         Returns
         -------
@@ -283,7 +306,7 @@ def split(self, n_parts: int, n_diffusing: int,
         """
         split_atom_sites = np.split(self.atom_sites, n_parts)
         split_transitions = _split_transitions_in_parts(
-            self.all_transitions, extras.n_steps, n_parts)
+            self.all_transitions, n_steps, n_parts)
 
         parts = [SitesData(self.structure) for _ in range(n_parts)]
 
@@ -300,7 +323,7 @@ def split(self, n_parts: int, n_diffusing: int,
             part.occupancy = part.calculate_occupancy()
 
             part.sites_occupancy = part.calculate_site_occupancy(
-                n_steps=extras.n_steps / n_parts)
+                n_steps=int(n_steps / n_parts))
 
             part.atom_locations = part.calculate_atom_locations(
                 n_diffusing=n_diffusing)
@@ -401,13 +424,13 @@ def calculate_jumps(self) -> dict[tuple[str, str], int]:
         return jumps
 
     def calculate_rates(
-            self, *, total_time: int,
+            self, *, total_time: float,
             n_diffusing: int) -> dict[tuple[str, str], tuple[float, float]]:
         """Calculate jump rates (total jumps / second).
 
         Parameters
         ----------
-        total_time : int
+        total_time : float
             Total time for the simulation
         n_diffusing : int
             Number of diffusing atoms
@@ -435,13 +458,13 @@ def calculate_rates(
         return rates
 
     def calculate_activation_energies(
-            self, *, total_time: int, n_diffusing: int, attempt_freq: float,
+            self, *, total_time: float, n_diffusing: int, attempt_freq: float,
             temperature: float) -> dict[tuple[str, str], tuple[float, float]]:
         """Calculate activation energies for jumps (UNITS?).
 
         Parameters
         ----------
-        total_time : int
+        total_time : float
             Total time for the simulation
         n_diffusing : int
             Number of diffusing atoms
diff --git a/src/trajectory.py b/src/trajectory.py
index d572cb0d..694b0622 100644
--- a/src/trajectory.py
+++ b/src/trajectory.py
@@ -36,14 +36,14 @@ class Trajectory(PymatgenTrajectory):
 
     def __init__(self, *, metadata: dict | None = None, **kwargs):
         super().__init__(**kwargs)
-        self.metadata = metadata if metadata else None
+        self.metadata = metadata if metadata else {}
 
     def __getitem__(self, frames):
         """Hack around pymatgen Trajectory limitations."""
         new = super().__getitem__(frames)
         if isinstance(new, PymatgenTrajectory):
             new.__class__ = self.__class__
-            new.metadata = self.metadata
+        new.metadata = self.metadata if hasattr(self, 'metadata') else {}
         return new
 
     def to_positions(self):
@@ -54,6 +54,16 @@ def to_positions(self):
         super().to_positions()
         self.coords = np.mod(self.coords, 1)
 
+    @property
+    def total_time(self) -> float:
+        """Return total time for trajectory."""
+        return len(self) * self.time_step
+
+    @property
+    def sampling_frequency(self) -> float:
+        """Return number of time steps per second."""
+        return 1 / self.time_step
+
     @property
     def positions(self) -> np.ndarray:
         """Return trajectory positions as fractional coordinates.
diff --git a/src/utils.py b/src/utils.py
index 3d24d8d5..38d18edd 100644
--- a/src/utils.py
+++ b/src/utils.py
@@ -1,5 +1,6 @@
 import numpy as np
 from pymatgen.core import Lattice, Structure
+from scipy import signal
 
 
 def ffill(arr: np.ndarray, fill_val: int = -1, axis=-1) -> np.ndarray:
@@ -44,6 +45,41 @@ def bfill(arr: np.ndarray, fill_val: int = -1, axis=-1) -> np.ndarray:
     return np.fliplr(ffill(np.fliplr(arr), fill_val=fill_val))
 
 
+def meanfreq(x: np.ndarray, fs: float = 1.0):
+    """Estimates the mean frequency in terms of the sample rate, fs.
+
+    Vectorized version of https://stackoverflow.com/a/56487241
+
+    Parameters
+    ----------
+    x : np.ndarray[i, j]
+        Time series of measurement values. The mean frequency is computed
+        along the last axis (-1).
+    fs : float, optional
+        Sampling frequency of the `x` time series. Defaults to 1.0.
+
+    Returns
+    -------
+    mnfreq : np.ndarray
+        Array of mean frequencies.
+    """
+    if x.ndim == 1:
+        x = x.reshape(1, -1)
+
+    assert x.ndim == 2
+
+    f, Pxx_den = signal.periodogram(x, fs, axis=-1)
+    width = np.tile(f[1] - f[0], Pxx_den.shape)
+    P = Pxx_den * width
+    pwr = np.sum(P, axis=1).reshape(-1, 1)
+
+    f = f.reshape(1, -1)
+
+    mnfreq = np.dot(P, f.T) / pwr
+
+    return mnfreq
+
+
 def is_lattice_similar(a: Lattice | Structure,
                        b: Lattice | Structure,
                        length_tol: float = 0.5,


Test patch:
diff --git a/tests/displacement_test.py b/tests/displacement_test.py
deleted file mode 100644
index 32a0992b..00000000
--- a/tests/displacement_test.py
+++ /dev/null
@@ -1,12 +0,0 @@
-import numpy as np
-
-
-def test_displacement_calculate_all(trajectory):
-    assert np.allclose(
-        trajectory.distances_from_base_position(),
-        np.array([[0., 0.2, 0.4, 0.6, 0.9], [0., 0., 0., 0., 0.]]))
-
-    diff_trajectory = trajectory.filter('B')
-
-    assert np.allclose(diff_trajectory.distances_from_base_position(),
-                       np.array([0., 0.2, 0.4, 0.6, 0.9]))
diff --git a/tests/integration_test.py b/tests/integration_test.py
index f62cce10..ee50ab9b 100644
--- a/tests/integration_test.py
+++ b/tests/integration_test.py
@@ -10,9 +10,9 @@
 import numpy as np
 import pytest
 from gemdat import SitesData
-from gemdat.calculate import calculate_all
 from gemdat.io import load_known_material
 from gemdat.rdf import calculate_rdfs
+from gemdat.simulation_metrics import SimulationMetrics
 from gemdat.trajectory import Trajectory
 from gemdat.volume import trajectory_to_volume
 
@@ -28,35 +28,19 @@
 
 
 @pytest.fixture
-def gemdat_results():
+def vasp_traj():
     trajectory = Trajectory.from_vasprun(VASP_XML)
     trajectory = trajectory[1250:]
-
-    extras = calculate_all(
-        trajectory,
-        diffusing_element='Li',
-        z_ion=1,
-        diffusion_dimensions=3,
-    )
-
-    return (trajectory, extras)
+    return trajectory
 
 
 @pytest.fixture
-def gemdat_results_subset():
+def vasp_traj_short():
     trajectory = Trajectory.from_vasprun(VASP_XML)
     # Reduced number of time steps for slow calculations
     trajectory = trajectory[4000:]
 
-    extras = calculate_all(
-        trajectory,
-        diffusing_element='Li',
-        z_ion=1,
-        diffusion_dimensions=3,
-        n_parts=1,
-    )
-
-    return (trajectory, extras)
+    return trajectory
 
 
 @pytest.fixture
@@ -65,10 +49,10 @@ def structure():
 
 
 @vaspxml_available
-def test_volume(gemdat_results):
-    trajectory, extras = gemdat_results
+def test_volume(vasp_traj):
+    trajectory = vasp_traj
 
-    diff_trajectory = trajectory.filter(extras.diffusing_element)
+    diff_trajectory = trajectory.filter('Li')
 
     vol = trajectory_to_volume(trajectory=diff_trajectory, resolution=0.2)
 
@@ -78,10 +62,10 @@ def test_volume(gemdat_results):
 
 
 @vaspxml_available
-def test_volume_cartesian(gemdat_results):
-    trajectory, extras = gemdat_results
+def test_volume_cartesian(vasp_traj):
+    trajectory = vasp_traj
 
-    diff_trajectory = trajectory.filter(extras.diffusing_element)
+    diff_trajectory = trajectory.filter('Li')
 
     vol = trajectory_to_volume(trajectory=diff_trajectory,
                                resolution=0.2,
@@ -93,25 +77,37 @@ def test_volume_cartesian(gemdat_results):
 
 
 @vaspxml_available
-def test_tracer(gemdat_results):
-    _, extras = gemdat_results
+def test_tracer(vasp_traj):
 
-    assert isclose(extras.particle_density, 2.4557e28, rel_tol=1e-4)
-    assert isclose(extras.mol_per_liter, 40.777, rel_tol=1e-4)
-    assert isclose(extras.tracer_diff, 1.5706e-09, rel_tol=1e-4)
-    assert isclose(extras.tracer_conduc, 110.322, rel_tol=1e-4)
+    diff_trajectory = vasp_traj.filter('Li')
+    metrics = SimulationMetrics(diff_trajectory)
+
+    assert isclose(metrics.particle_density(), 2.4557e28, rel_tol=1e-4)
+    assert isclose(metrics.mol_per_liter(), 40.777, rel_tol=1e-4)
+    assert isclose(metrics.tracer_diffusivity(diffusion_dimensions=3),
+                   1.5706e-09,
+                   rel_tol=1e-4)
+    assert isclose(metrics.tracer_conductivity(z_ion=1,
+                                               diffusion_dimensions=3),
+                   110.322,
+                   rel_tol=1e-4)
 
 
 @vaspxml_available
-def test_sites(gemdat_results, structure):
-    trajectory, extras = gemdat_results
+def test_sites(vasp_traj, structure):
+    trajectory = vasp_traj
 
-    sites = SitesData(structure)
-    sites.calculate_all(trajectory=trajectory, extras=extras)
+    n_parts = 10
 
-    n_steps = extras.n_steps
-    n_diffusing = sum(
-        [sp.symbol == extras.diffusing_element for sp in trajectory.species])
+    sites = SitesData(structure)
+    sites.calculate_all(trajectory=trajectory,
+                        diffusing_element='Li',
+                        z_ion=1,
+                        diffusion_dimensions=3,
+                        n_parts=n_parts)
+
+    n_steps = len(trajectory)
+    n_diffusing = sum([sp.symbol == 'Li' for sp in trajectory.species])
     n_sites = sites.n_sites
 
     assert sites.atom_sites.shape == (n_steps, n_diffusing)
@@ -125,14 +121,14 @@ def test_sites(gemdat_results, structure):
 
     assert sites.transitions.shape == (n_sites, n_sites)
 
-    assert sites.transitions_parts.shape == (extras.n_parts, n_sites, n_sites)
+    assert sites.transitions_parts.shape == (n_parts, n_sites, n_sites)
     assert np.sum(sites.transitions_parts[0]) == 37
     assert np.sum(sites.transitions_parts[9]) == 38
 
     assert sites.occupancy[0] == 1704
     assert sites.occupancy[43] == 542
 
-    assert len(sites.occupancy_parts) == extras.n_parts
+    assert len(sites.occupancy_parts) == n_parts
 
     assert sites.occupancy_parts[0][0] == 56
     assert sites.occupancy_parts[0][42] == 36
@@ -141,7 +137,7 @@ def test_sites(gemdat_results, structure):
 
     assert isclose(sites.sites_occupancy['48h'], 0.380628, rel_tol=1e-4)
 
-    assert len(sites.sites_occupancy_parts) == extras.n_parts
+    assert len(sites.sites_occupancy_parts) == n_parts
     assert isclose(sites.sites_occupancy_parts[0]['48h'],
                    0.37756,
                    rel_tol=1e-4)
@@ -151,7 +147,7 @@ def test_sites(gemdat_results, structure):
 
     assert isclose(sites.atom_locations['48h'], 0.761255, rel_tol=1e-4)
 
-    assert len(sites.atom_locations_parts) == extras.n_parts
+    assert len(sites.atom_locations_parts) == n_parts
     assert isclose(sites.atom_locations_parts[0]['48h'],
                    0.755111,
                    rel_tol=1e-4)
@@ -198,18 +194,24 @@ def test_sites(gemdat_results, structure):
 
 
 @vaspxml_available
-def test_rdf(gemdat_results_subset, structure):
-    trajectory, extras = gemdat_results_subset
+def test_rdf(vasp_traj_short, structure):
+    trajectory = vasp_traj_short
 
     structure = load_known_material('argyrodite')
 
     sites = SitesData(structure)
-    sites.calculate_all(trajectory=trajectory, extras=extras)
+    sites.calculate_all(
+        trajectory=trajectory,
+        diffusing_element='Li',
+        z_ion=1,
+        diffusion_dimensions=3,
+        n_parts=1,
+    )
 
     rdfs = calculate_rdfs(
         trajectory=trajectory,
         sites=sites,
-        species=extras.diffusing_element,
+        species='Li',
         max_dist=5,
     )
 
diff --git a/tests/simulation_metrics_test.py b/tests/simulation_metrics_test.py
new file mode 100644
index 00000000..0bcb19e1
--- /dev/null
+++ b/tests/simulation_metrics_test.py
@@ -0,0 +1,29 @@
+import numpy as np
+from gemdat.simulation_metrics import SimulationMetrics
+
+
+def test_tracer_metrics(trajectory):
+    diff_trajectory = trajectory.filter('B')
+    metrics = SimulationMetrics(diff_trajectory)
+
+    assert (np.isclose(metrics.particle_density(), 1e30))
+    assert (np.isclose(metrics.mol_per_liter(), 1660.53906))
+    assert (np.isclose(metrics.tracer_diffusivity(diffusion_dimensions=3),
+                       2.666667e-10))
+    assert (np.isclose(
+        metrics.tracer_conductivity(z_ion=1, diffusion_dimensions=3),
+        4.081278e-09))
+
+
+def test_vibration_metrics(trajectory):
+    diff_trajectory = trajectory.filter('B')
+    metrics = SimulationMetrics(diff_trajectory)
+
+    assert np.allclose(metrics.speed(), [[0., 0.2, 0.2, 0.2, 0.3]])
+
+    attempt_freq, attempt_freq_std = metrics.attempt_frequency()
+
+    assert np.isclose(attempt_freq, 0.33727)
+    assert np.isclose(attempt_freq_std, 0.0)
+    assert np.allclose(metrics.amplitudes(), np.array([0.9]))
+    assert np.isclose(metrics.vibration_amplitude(), 0.0)
diff --git a/tests/tracer_test.py b/tests/tracer_test.py
deleted file mode 100644
index feb196f0..00000000
--- a/tests/tracer_test.py
+++ /dev/null
@@ -1,21 +0,0 @@
-from types import SimpleNamespace
-
-import numpy as np
-from gemdat.calculate.tracer import Tracer
-
-
-def test_vibration_calculate_all(trajectory):
-    extras = SimpleNamespace(
-        total_time=1e-12,
-        n_diffusing=1,
-        diffusion_dimensions=3,
-        z_ion=1,
-    )
-    diff_trajectory = trajectory.filter('B')
-
-    ret = Tracer.calculate_all(diff_trajectory, extras)
-
-    assert (np.isclose(ret['particle_density'], 1e30))
-    assert (np.isclose(ret['mol_per_liter'], 1660.53906))
-    assert (np.isclose(ret['tracer_diff'], 2.666667e-10))
-    assert (np.isclose(ret['tracer_conduc'], 20406.389030))
diff --git a/tests/utils_test.py b/tests/utils_test.py
index 05bf189b..1a954758 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-from gemdat.utils import bfill, ffill
+from gemdat.utils import bfill, ffill, meanfreq
 
 
 @pytest.fixture
@@ -54,3 +54,25 @@ def test_bfill_axis0(arr):
     ])
 
     np.testing.assert_equal(ret, expected)
+
+
+def test_meanfreq_single_timestep():
+    x = np.sin(np.linspace(0, 1, 6))
+    ret = meanfreq(x)
+
+    expected = np.array([[0.2303359]])
+
+    np.testing.assert_allclose(ret, expected)
+
+
+def test_meanfreq():
+    x = np.array([
+        np.sin(np.linspace(0, 1, 6)),
+        np.sin(np.linspace(0, 2, 6)),
+        np.sin(np.linspace(0, 3, 6)),
+    ])
+    ret = meanfreq(x)
+
+    expected = np.array([[0.2303359], [0.21308077], [0.17074241]])
+
+    np.testing.assert_allclose(ret, expected)
diff --git a/tests/vibration_test.py b/tests/vibration_test.py
deleted file mode 100644
index a86e3515..00000000
--- a/tests/vibration_test.py
+++ /dev/null
@@ -1,35 +0,0 @@
-import numpy as np
-from gemdat.calculate import Vibration
-from gemdat.calculate.vibration import meanfreq
-
-
-def test_vibration_calculate_all(trajectory):
-    ret = Vibration.calculate_all(trajectory.filter('B'))
-
-    assert np.allclose(ret['speed'], np.array([[0., 0.2, 0.2, 0.2, 0.3]]))
-    assert np.isclose(ret['attempt_freq'], 0.33727)
-    assert np.isclose(ret['attempt_freq_std'], 0.0)
-    assert np.allclose(ret['amplitudes'], np.array([0.9]))
-    assert np.isclose(ret['vibration_amplitude'], 0.0)
-
-
-def test_meanfreq_single_timestep():
-    x = np.sin(np.linspace(0, 1, 6))
-    ret = meanfreq(x)
-
-    expected = np.array([[0.2303359]])
-
-    np.testing.assert_allclose(ret, expected)
-
-
-def test_meanfreq():
-    x = np.array([
-        np.sin(np.linspace(0, 1, 6)),
-        np.sin(np.linspace(0, 2, 6)),
-        np.sin(np.linspace(0, 3, 6)),
-    ])
-    ret = meanfreq(x)
-
-    expected = np.array([[0.2303359], [0.21308077], [0.17074241]])
-
-    np.testing.assert_allclose(ret, expected)