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

Issue:
Add some examples to the documentation



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index f8d3aa7a..febc1877 100644
--- a/README.md
+++ b/README.md
@@ -4,21 +4,25 @@
 <!-- [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/XXX)](https://pypi.org/project/XXX/) -->
 <!-- [![PyPI](https://img.shields.io/pypi/v/XXX.svg?style=flat)](https://pypi.org/project/XXX/) -->
 
-![GEMDAT banner](https://raw.githubusercontent.com/GEMDAT-repos/GEMDAT/main/src/data/logo.png)
+![GEMDAT banner](https://raw.githubusercontent.com/GEMDAT-repos/GEMDAT/main/src/gemdat/data/logo.png)
 
 # GEMDAT
 
-This repository contains Python code to analyse Molecular Dynamics simulations.
+Gemdat is a Python library for the analysis of diffusion in solid-state electrolytes from Molecular Dynamics simulations. Gemdat is built on top of [Pymatgen](https://pymatgen.org/), making it straightforward to integrate it into your Pymatgen-based workflows.
 
-The code in this repository is based on:
-https://bitbucket.org/niekdeklerk/md-analysis-with-matlab/src/master/
+With Gemdat, you can:
 
-## Installation
+- Explore your MD simulation via an easy-to-use Python API
+- Load and analyze trajectories from VASP simulation data
+- Characterize and visualize diffusivity
+- Find jumps and transitions between sites
+- Plot radial distribution functions
+- Analyze results via an interactive dashboard
 
 To install:
 
 ```console
-pip install .[develop]
+pip install gemdat
 ```
 
 The source code is available from [Github](https://github.com/GEMDAT-repos/GEMDAT).
@@ -35,13 +39,14 @@ from gemdat.calculate.extras import calculate_all
 
 trajectory = Trajectory.from_vasprun(Path('../example/vasprun.xml'))
 
+plots.plot_displacement_per_element(trajectory=trajectory)
+
 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)
+plots.plot_displacement_per_site(trajectory=diff_trajectory)
+plots.plot_displacement_histogram(trajectory=diff_trajectory)
+plots.plot_frequency_vs_occurence(trajectory=diff_trajectory)
+plots.plot_vibrational_amplitudes(trajectory=diff_trajectory)
 
 structure = load_known_material('argyrodite', supercell=(2, 1, 1))
 
@@ -51,10 +56,10 @@ sites = SitesData(
    floating_specie='Li',
 )
 
-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)
+plots.plot_jumps_vs_distance(sites=sites)
+plots.plot_jumps_vs_time(sites=sites)
+plots.plot_collective_jumps(sites=sites)
+plots.plot_jumps_3d(sites=sites)
 ```
 
 Or, one function to do everything:
@@ -76,4 +81,8 @@ Check out our [Contributing Guidelines](CONTRIBUTING.md#Getting-started-with-dev
 
 ## References
 
+The code in this repository is based on [Matlab code to analyse Molecular Dynamics simulations](https://bitbucket.org/niekdeklerk/md-analysis-with-matlab/src/master/).
+
+For background information on how some of the properties are calculated, check out the accompanying paper:
+
 - Niek J.J. de Klerk, Eveline van der Maas and Marnix Wagemaker, ACS Applied Energy Materials, (2018), doi: [10.1021/acsaem.8b00457](https://doi.org/10.1021/acsaem.8b00457)
diff --git a/docs/api.md b/docs/api.md
deleted file mode 100644
index 247f6ad2..00000000
--- a/docs/api.md
+++ /dev/null
@@ -1,1 +0,0 @@
-::: gemdat
diff --git a/docs/api/gemdat.md b/docs/api/gemdat.md
new file mode 100644
index 00000000..493ce1ce
--- /dev/null
+++ b/docs/api/gemdat.md
@@ -0,0 +1,8 @@
+- [gemdat.load_cif][gemdat.io.load_cif]
+- [gemdat.load_known_material][gemdat.io.load_known_material]
+- [gemdat.SimulationMetrics][gemdat.simulation_metrics.SimulationMetrics]
+- [gemdat.SitesData][gemdat.sites.SitesData]
+- [gemdat.Trajectory][gemdat.trajectory.Trajectory]
+- [gemdat.trajectory_to_vasp_volume][gemdat.volume.trajectory_to_vasp_volume]
+- [gemdat.trajectory_to_volume][gemdat.volume.trajectory_to_volume]
+- [gemdat.volume_to_structure][gemdat.volume.volume_to_structure]
diff --git a/docs/api/gemdat_io.md b/docs/api/gemdat_io.md
new file mode 100644
index 00000000..dcb81b64
--- /dev/null
+++ b/docs/api/gemdat_io.md
@@ -0,0 +1,4 @@
+::: gemdat.io
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false
diff --git a/docs/api/gemdat_legacy.md b/docs/api/gemdat_legacy.md
new file mode 100644
index 00000000..7aff5582
--- /dev/null
+++ b/docs/api/gemdat_legacy.md
@@ -0,0 +1,4 @@
+::: gemdat.legacy
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false
diff --git a/docs/api/gemdat_plots.md b/docs/api/gemdat_plots.md
new file mode 100644
index 00000000..d1110cf5
--- /dev/null
+++ b/docs/api/gemdat_plots.md
@@ -0,0 +1,4 @@
+::: gemdat.plots
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false
diff --git a/docs/api/gemdat_rdf.md b/docs/api/gemdat_rdf.md
new file mode 100644
index 00000000..6f5de139
--- /dev/null
+++ b/docs/api/gemdat_rdf.md
@@ -0,0 +1,4 @@
+::: gemdat.rdf
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false
diff --git a/docs/api/gemdat_simulation_metrics.md b/docs/api/gemdat_simulation_metrics.md
new file mode 100644
index 00000000..3bd2eb19
--- /dev/null
+++ b/docs/api/gemdat_simulation_metrics.md
@@ -0,0 +1,4 @@
+::: gemdat.simulation_metrics
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false
diff --git a/docs/api/gemdat_sites.md b/docs/api/gemdat_sites.md
new file mode 100644
index 00000000..70eaa72d
--- /dev/null
+++ b/docs/api/gemdat_sites.md
@@ -0,0 +1,5 @@
+::: gemdat.sites
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false
+      merge_init_into_class: true
diff --git a/docs/api/gemdat_trajectory.md b/docs/api/gemdat_trajectory.md
new file mode 100644
index 00000000..a4a570ba
--- /dev/null
+++ b/docs/api/gemdat_trajectory.md
@@ -0,0 +1,5 @@
+::: gemdat.trajectory
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false
+      merge_init_into_class: true
diff --git a/docs/api/gemdat_utils.md b/docs/api/gemdat_utils.md
new file mode 100644
index 00000000..df3bf288
--- /dev/null
+++ b/docs/api/gemdat_utils.md
@@ -0,0 +1,4 @@
+::: gemdat.utils
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false
diff --git a/docs/api/gemdat_volume.md b/docs/api/gemdat_volume.md
new file mode 100644
index 00000000..e1d23da2
--- /dev/null
+++ b/docs/api/gemdat_volume.md
@@ -0,0 +1,4 @@
+::: gemdat.volume
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false
diff --git a/docs/dashboard.md b/docs/dashboard.md
new file mode 100644
index 00000000..d76164ee
--- /dev/null
+++ b/docs/dashboard.md
@@ -0,0 +1,23 @@
+Gemdash includes a dashboard to interactively select your input data and experimental parameters. This makes it straightforward to explore your data without any coding experience required.
+
+To use the dashboard, first make sure you install all the dependencies:
+
+```
+pip install gemdat[dashboard]
+```
+
+The dashboard is available via the command
+
+```
+gemdash
+```
+
+To pre-set your data, you can set the `VASP_XML` environment variable (useful for testing):
+
+```
+VASP_XML=./path/to/vasprun.xml gemdash
+```
+
+!!! note
+
+    The dashboard runs in the web browser and needs a somewhat recent version of [firefox or one of the other supported browsers here](https://docs.streamlit.io/knowledge-base/using-streamlit/supported-browsers).
diff --git a/docs/getting_started.md b/docs/getting_started.md
new file mode 100644
index 00000000..c46a2b18
--- /dev/null
+++ b/docs/getting_started.md
@@ -0,0 +1,82 @@
+This page shows you how to get started using Gemdat in your own Python scripts or in a Jupyter notebook.
+
+## Loading simulation data
+
+The first step is to load your simulation data. These are represented in gemdat by a trajectory. A trajectory is essentially a list of coordinates for all species over time. To load a trajectory:
+
+```python
+from gemdat import Trajectory
+
+trajectory = Trajectory.from_vasprun('./path/to/vasprun.xml')
+```
+
+For now, select the trajectory for the diffusing species, Lithium, by its symbol:
+
+```python
+diff_trajectory = trajectory.filter('Li')
+```
+
+## Plotting trajectory properties
+
+You can do this by importing the `plots` submodule. In the example below, we are plotting the vibration of the diffusing element by passing its trajectory.
+
+```python
+from gemdat import plots
+
+plots.plot_vibrational_amplitudes(diff_trajectory)
+```
+
+See the [plots page](./plotting.md#trajectory-and-displacements-plots) for more information about which plots for trajectories are available.
+
+## Analyzing jumps and transitions
+
+To find out how Lithium jumps from site to site, we must know what the sites are. Gemdat contains a [small internal database](https://github.com/GEMDAT-repos/GEMDAT/tree/main/src/gemdat/data) with available structures.
+
+Alternatively, you can load the sites from 1. [a cif file][gemdat.io.load_cif], 2. [a density or volume][gemdat.volume.volume_to_structure] or 3. construct your own [Structure][pymatgen.core.structure.Structure].
+
+```python
+from gemdat import load_known_material
+
+structure = load_known_material('argyrodite', supercell=(2, 1, 1))
+```
+
+The [gemdat.SitesData][] class is responsible for calculating the different jumps properties. We pass the `structure` with the jump sites, the trajectory containing both the host and diffusing species, and specify what the diffusing element is.
+
+```python
+from gemdat import SitesData
+
+sites = SitesData(
+   structure=structure,
+   trajectory=trajectory,
+   floating_specie='Li',
+)
+```
+
+See the documentation for [`gemdat.SitesData`][] to find out which properties are available.
+
+## Plotting jumps and transition properties
+
+Gemdat contains [several functions for plotting jumps and transitions](./plotting.md#jumps-and-transition-plots). These take the [gemdat.SitesData][] object we just constructed as input.
+
+For example, to visualize all jumps in 3d:
+
+```
+plots.plot_jumps_3d(sites)
+```
+
+## Radial distribution functions
+
+This function calculates the [radial distribution function](https://en.wikipedia.org/wiki/Radial_distribution_function) (RDF) from the `floating_specie` on [gemdat.SitesData][],
+
+```python
+from gemdat.rdf import radial_distribution
+
+rdf_data = radial_distribution(sites=sites)
+```
+
+This returns a dictionary with all possible transitions between the jump sites. You can use the [plotting module](./plotting.md#radial-distribution-plots) to visualize the RDFs:
+
+```python
+for rdfs in rdf_data.values():
+   plots.radial_distribution(rdfs)
+```
diff --git a/docs/plotting.md b/docs/plotting.md
new file mode 100644
index 00000000..8d9e1d1b
--- /dev/null
+++ b/docs/plotting.md
@@ -0,0 +1,60 @@
+# Plots
+
+Gemdat contains several built-in plots for visualizing trajectories, jumps, transitions, and radial distribution functions.
+
+These are collected in the `plots` module. The intended usage is that you import `gemdat.plots` like this:
+
+```python
+from gemdat import plots
+
+plots.displacement_per_element(trajectory)
+plots.jumps_vs_distance(trajectory, sites)
+plots.radial_distribution(rdfs)
+```
+
+All plotting functions take a [gemdat.Trajectory][], [gemdat.SitesData][], [gemdat.rdf.RDFData][] or a combination as input. In addition, for some plots you have a few parameters to tune the output.
+
+
+## Trajectory and displacements plots
+
+::: gemdat.plots
+    options:
+      show_root_toc_entry: false
+      heading_level: 3
+      members:
+        - displacement_per_site
+        - displacement_per_element
+        - displacement_histogram
+
+
+## Simulation metrics plots
+
+::: gemdat.plots
+    options:
+      show_root_toc_entry: false
+      heading_level: 3
+      members:
+        - frequency_vs_occurence
+        - vibrational_amplitudes
+
+## Jumps and transition plots
+
+::: gemdat.plots
+    options:
+      show_root_toc_entry: false
+      heading_level: 3
+      members:
+        - jumps_vs_distance
+        - jumps_vs_time
+        - collective_jumps
+        - jumps_3d
+        - jumps_3d_animation
+
+## Radial distribution plots
+
+::: gemdat.plots
+    options:
+      show_root_toc_entry: false
+      heading_level: 3
+      members:
+        - radial_distribution
diff --git a/docs/stylesheets/extra.css b/docs/stylesheets/extra.css
new file mode 100644
index 00000000..e69de29b
diff --git a/mkdocs.yml b/mkdocs.yml
index b6fe1e48..997620ec 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -5,13 +5,25 @@ repo_name: GitHub
 
 nav:
   - Home: index.md
-  - Python API: api.md
+  - Getting started: getting_started.md
+  - Visualization: plotting.md
+  - Dashboard: dashboard.md
+  - Python API:
+    - gemdat: api/gemdat.md
+    - gemdat.io: api/gemdat_io.md
+    - gemdat.legacy: api/gemdat_legacy.md
+    - gemdat.plots: api/gemdat_plots.md
+    - gemdat.rdf: api/gemdat_rdf.md
+    - gemdat.simulation_metrics: api/gemdat_simulation_metrics.md
+    - gemdat.sites: api/gemdat_sites.md
+    - gemdat.trajectory: api/gemdat_trajectory.md
+    - gemdat.utils: api/gemdat_utils.md
+    - gemdat.volume: api/gemdat_volume.md
   - Contributing: CONTRIBUTING.md
   - Code of Conduct: CODE_OF_CONDUCT.md
   - 🔗 Source code: https://github.com/GEMDAT-repos/GEMDAT
   - 🔗 Issues: https://github.com/GEMDAT-repos/GEMDAT/issues
 
-
 theme:
   name: material
   primary: blue
@@ -42,20 +54,25 @@ plugins:
     # https://mkdocstrings.github.io/usage/
     handlers:
       python:
+        paths: [src]
         import:
         - https://docs.python.org/3/objects.inv
         - https://numpy.org/doc/stable/objects.inv
         - https://docs.scipy.org/doc/scipy/objects.inv
         - https://pandas.pydata.org/docs/objects.inv
+        - https://matplotlib.org/stable/objects.inv
+        - https://pymatgen.org/objects.inv
+        - https://scikit-image.org/docs/stable/objects.inv
         options:
           docstring_style: numpy
           docstring_options:
             ignore_init_summary: yes
           show_submodules: no
           show_source: true
+          show_root_heading: false
+          show_root_full_path: false
           docstring_section_style: list
           members_order: alphabetical
-          show_category_heading: no
           merge_init_into_class: yes
           filters: ["!^_"]
 
diff --git a/src/gemdat/collective.py b/src/gemdat/collective.py
index c12c4227..02f841bc 100644
--- a/src/gemdat/collective.py
+++ b/src/gemdat/collective.py
@@ -32,9 +32,9 @@ def __init__(self,
         ----------
         transitions : Transitions
             Input transition events
-        structure : Structure
+        structure : pymatgen.core.structure.Structure
             Structure with list of jump sites
-        lattice : Lattice
+        lattice : pymatgen.core.lattice.Lattice
             Input lattice for distance calculations (from simulation data)
         max_steps : int
             Maximum number of time steps which would still mean correlation
diff --git a/src/gemdat/dashboard/run.py b/src/gemdat/dashboard/run.py
index e7298dc2..903758d4 100644
--- a/src/gemdat/dashboard/run.py
+++ b/src/gemdat/dashboard/run.py
@@ -164,10 +164,10 @@ def _SitesData(**kwargs):
         plots.displacement_histogram(trajectory=diff_trajectory),
         plots.frequency_vs_occurence(trajectory=diff_trajectory),
         plots.vibrational_amplitudes(trajectory=diff_trajectory),
-        plots.jumps_vs_distance(trajectory=trajectory, sites=sites),
-        plots.jumps_vs_time(trajectory=trajectory, sites=sites),
-        plots.collective_jumps(trajectory=trajectory, sites=sites),
-        plots.jumps_3d(trajectory=trajectory, sites=sites),
+        plots.jumps_vs_distance(sites=sites),
+        plots.jumps_vs_time(sites=sites),
+        plots.collective_jumps(sites=sites),
+        plots.jumps_3d(sites=sites),
     )
 
     # automagically divide the plots over the number of columns
@@ -197,8 +197,7 @@ def _radial_distribution(**kwargs):
             )
 
         rdf_figures = [
-            plots.radial_distribution(rdfs.values())
-            for rdfs in rdf_data.values()
+            plots.radial_distribution(rdfs) for rdfs in rdf_data.values()
         ]
 
         # automagically divide the plots over the number of columns
diff --git a/src/gemdat/io.py b/src/gemdat/io.py
index 42bec82d..a36fe17a 100644
--- a/src/gemdat/io.py
+++ b/src/gemdat/io.py
@@ -1,3 +1,5 @@
+"""The io module contains functions to read and write data."""
+
 from importlib.resources import files
 from pathlib import Path
 
@@ -18,7 +20,8 @@
 
 
 def load_cif(filename: Path | str) -> Structure:
-    """Load cif file and return first item as pymatgen structure.
+    """Load cif file and return first item as
+    [pymatgen.core.structure.Structure][].
 
     Parameters
     ----------
@@ -27,7 +30,7 @@ def load_cif(filename: Path | str) -> Structure:
 
     Returns
     -------
-    structure : Structure
+    structure : pymatgen.core.structure.Structure
         Output structure
     """
     cifdata = cif.CifParser(filename)
@@ -35,9 +38,10 @@ def load_cif(filename: Path | str) -> Structure:
     return structure
 
 
-def load_known_material(name: str,
-                        supercell: tuple[int, int, int] | None = None
-                        ) -> Structure:
+def load_known_material(
+    name: str,
+    supercell: tuple[int, int, int] | None = None,
+) -> Structure:
     """Load known material from internal database.
 
     Parameters
@@ -46,12 +50,12 @@ def load_known_material(name: str,
         Name of the material
     supercell : tuple(int, int, int) | None, optional
         Optionally, scale the lattice by a sequence of three factors.
-        For example, (2, 1, 1) specifies that the supercell should have
-        dimensions 2a x b x c.
+        For example, `(2, 1, 1)` specifies that the supercell should have
+        dimensions $2a \\times b \\times c$.
 
     Returns
     -------
-    structure : Structure
+    structure : pymatgen.core.structure.Structure
         Output structure
     """
     filename = (DATA / name).with_suffix('.cif')
diff --git a/src/gemdat/legacy.py b/src/gemdat/legacy.py
index 858aceb5..007a873e 100644
--- a/src/gemdat/legacy.py
+++ b/src/gemdat/legacy.py
@@ -1,3 +1,6 @@
+"""This module provides a code to deal with the original matlab code that
+Gemdat is based on."""
+
 from gemdat import SitesData, load_known_material, plots
 from gemdat.rdf import radial_distribution
 from gemdat.trajectory import Trajectory
@@ -25,6 +28,10 @@ def analyse_md(
 ) -> tuple[Trajectory, SitesData]:
     """Analyse md data.
 
+    This function mimicks the the API of the `analyse_md` function in the
+    [Matlab code to analyse Molecular Dynamics simulations](https://bitbucket.org/niekdeklerk/md-analysis-with-matlab/src/master/)
+    that Gemdat is based on.
+
     Parameters
     ----------
     vasp_xml : str
@@ -62,8 +69,10 @@ def analyse_md(
 
     Returns
     -------
-    tuple[SimulationData, SitesData, SimpleNamespace]
-        Output data
+    trajectory : Trajectory
+        Output trajectory
+    sites : SitesData
+        Output sites data
     """
     trajectory = Trajectory.from_vasprun(vasp_xml)
 
@@ -87,15 +96,12 @@ def analyse_md(
     plots.displacement_histogram(trajectory=diff_trajectory)
     plots.frequency_vs_occurence(trajectory=diff_trajectory)
     plots.vibrational_amplitudes(trajectory=diff_trajectory)
-    plots.jumps_vs_distance(trajectory=trajectory,
-                            sites=sites,
-                            jump_res=jump_res)
-    plots.jumps_vs_time(trajectory=trajectory, sites=sites)
-    plots.collective_jumps(trajectory=trajectory, sites=sites)
-    plots.jumps_3d(trajectory=trajectory, sites=sites)
+    plots.jumps_vs_distance(sites=sites, jump_res=jump_res)
+    plots.jumps_vs_time(sites=sites)
+    plots.collective_jumps(sites=sites)
+    plots.jumps_3d(sites=sites)
 
     plots.jumps_3d_animation(
-        trajectory=trajectory,
         sites=sites,
         t_start=start_end[0],
         t_stop=start_end[1],
@@ -119,6 +125,6 @@ def analyse_md(
             resolution=rdf_res,
         )
         for rdfs in rdf_data.values():
-            plots.radial_distribution(rdfs.values())
+            plots.radial_distribution(rdfs)
 
     return trajectory, sites
diff --git a/src/gemdat/plots/__init__.py b/src/gemdat/plots/__init__.py
index 86e35f1a..5179726a 100644
--- a/src/gemdat/plots/__init__.py
+++ b/src/gemdat/plots/__init__.py
@@ -1,3 +1,5 @@
+"""This module contains all the plots that Gemdat can generate."""
+
 from ._displacements import (
     displacement_histogram,
     displacement_per_element,
diff --git a/src/gemdat/plots/_displacements.py b/src/gemdat/plots/_displacements.py
index dd3a64c4..7cda4f15 100644
--- a/src/gemdat/plots/_displacements.py
+++ b/src/gemdat/plots/_displacements.py
@@ -10,12 +10,12 @@ def displacement_per_site(*, trajectory: Trajectory) -> plt.Figure:
     Parameters
     ----------
     trajectory : Trajectory
-        Input trajectory
+        Input trajectory, i.e. for the diffusing atom
 
     Returns
     -------
-    fig : plt.Figure
-        Output matplotlib figure
+    fig : matplotlib.figure.Figure
+        Output figure
     """
     fig, ax = plt.subplots()
 
@@ -39,8 +39,8 @@ def displacement_per_element(*, trajectory: Trajectory) -> plt.Figure:
 
     Returns
     -------
-    fig : plt.Figure
-        Output matplotlib figure
+    fig : matplotlib.figure.Figure
+        Output figure
     """
     from collections import defaultdict
 
@@ -67,18 +67,17 @@ def displacement_per_element(*, trajectory: Trajectory) -> plt.Figure:
 
 
 def displacement_histogram(trajectory: Trajectory) -> plt.Figure:
-    """Plot histogram of total displacement of diffusing element at final
-    timestep.
+    """Plot histogram of total displacement at final timestep.
 
     Parameters
     ----------
     trajectory : Trajectory
-        Input trajectory
+        Input trajectory, i.e. for the diffusing atom
 
     Returns
     -------
-    fig : plt.Figure
-        Output matplotlib figure
+    fig : matplotlib.figure.Figure
+        Output figure
     """
     fig, ax = plt.subplots()
     ax.hist(trajectory.distances_from_base_position()[:, -1])
diff --git a/src/gemdat/plots/_jumps.py b/src/gemdat/plots/_jumps.py
index ff36debb..d4403aa3 100644
--- a/src/gemdat/plots/_jumps.py
+++ b/src/gemdat/plots/_jumps.py
@@ -10,19 +10,17 @@
 
 if TYPE_CHECKING:
     from gemdat import SitesData
-    from gemdat.trajectory import Trajectory
 
 
-def jumps_vs_distance(*,
-                      trajectory: Trajectory,
-                      sites: SitesData,
-                      jump_res: float = 0.1) -> plt.Figure:
+def jumps_vs_distance(
+    *,
+    sites: SitesData,
+    jump_res: float = 0.1,
+) -> plt.Figure:
     """Plot jumps vs. distance histogram.
 
     Parameters
     ----------
-    trajectory : Trajectory
-        Input trajectory
     sites : SitesData
         Input sites data
     jump_res : float, optional
@@ -30,10 +28,14 @@ def jumps_vs_distance(*,
 
     Returns
     -------
-    plt.Figure
+    fig : matplotlib.figure.Figure
+        Output figure
     """
-    lattice = trajectory.get_lattice()
     structure = sites.structure
+
+    trajectory = sites.trajectory
+    lattice = trajectory.get_lattice()
+
     pdist = lattice.get_all_distances(structure.frac_coords,
                                       structure.frac_coords)
 
@@ -57,16 +59,11 @@ def jumps_vs_distance(*,
     return fig
 
 
-def jumps_vs_time(*,
-                  trajectory: Trajectory,
-                  sites: SitesData,
-                  binsize: int = 500) -> plt.Figure:
+def jumps_vs_time(*, sites: SitesData, binsize: int = 500) -> plt.Figure:
     """Plot jumps vs. distance histogram.
 
     Parameters
     ----------
-    trajectory : Trajectory
-        Input trajectory
     sites : SitesData
         Input sites data
     binsize : int, optional
@@ -74,8 +71,11 @@ def jumps_vs_time(*,
 
     Returns
     -------
-    plt.Figure
+    fig : matplotlib.figure.Figure
+        Output figure
     """
+    trajectory = sites.trajectory
+
     n_steps = len(trajectory)
     bins = np.arange(0, n_steps + binsize, binsize)
 
@@ -90,20 +90,18 @@ def jumps_vs_time(*,
     return fig
 
 
-def collective_jumps(*, trajectory: Trajectory,
-                     sites: SitesData) -> plt.Figure:
+def collective_jumps(*, sites: SitesData) -> plt.Figure:
     """Plot collective jumps per jump-type combination.
 
     Parameters
     ----------
-    trajectory : Trajectory
-        Input trajectory
     sites : SitesData
         Input sites data
 
     Returns
     -------
-    plt.Figure
+    fig : matplotlib.figure.Figure
+        Output figure
     """
     fig, ax = plt.subplots()
 
@@ -121,20 +119,20 @@ def collective_jumps(*, trajectory: Trajectory,
     return fig
 
 
-def jumps_3d(*, trajectory: Trajectory, sites: SitesData) -> plt.Figure:
+def jumps_3d(*, sites: SitesData) -> plt.Figure:
     """Plot jumps in 3D.
 
     Parameters
     ----------
-    trajectory : Trajectory
-        Input trajectory
     sites : SitesData
         Input sites data
 
     Returns
     -------
-    plt.Figure
+    fig : matplotlib.figure.Figure
+        Output figure
     """
+    trajectory = sites.trajectory
 
     class LabelItems:
 
@@ -205,26 +203,19 @@ def items(self):
     return fig
 
 
-def jumps_3d_animation(*,
-                       trajectory: Trajectory,
-                       sites: SitesData,
-                       t_start: int,
-                       t_stop: int,
-                       decay: float = 0.05,
-                       skip: int = 5,
-                       interval: int = 20):
+def jumps_3d_animation(
+    *,
+    sites: SitesData,
+    t_start: int,
+    t_stop: int,
+    decay: float = 0.05,
+    skip: int = 5,
+    interval: int = 20,
+) -> plt.Figure:
     """Plot jumps in 3D as an animation over time.
 
-    # TODO
-    # - Refactor using init func
-    # - Refactor shared code with `plot_jumps_3d`
-    # - Save/export animation somehow
-    # - Probably cleaner to combine these functions as a class
-
     Parameters
     ----------
-    trajectory : Trajectory
-        Input trajectory
     sites : SitesData
         Input sites data
     t_start : int
@@ -240,11 +231,14 @@ def jumps_3d_animation(*,
 
     Returns
     -------
-    plt.Figure
+    fig : matplotlib.figure.Figure
+        Output figure
     """
     minwidth = 0.2
     maxwidth = 5.0
 
+    trajectory = sites.trajectory
+
     class LabelItems:
 
         def __init__(self, labels, coords):
diff --git a/src/gemdat/plots/_rdf.py b/src/gemdat/plots/_rdf.py
index ea1e6f24..22564e7f 100644
--- a/src/gemdat/plots/_rdf.py
+++ b/src/gemdat/plots/_rdf.py
@@ -13,13 +13,13 @@ def radial_distribution(rdfs: Iterable[RDFData]) -> plt.Figure:
 
     Parameters
     ----------
-    rdfs : dict[str, np.ndarray]
-        Dictionary with rdf array per symbol
+    rdfs : Iterable[RDFData]
+        List of RDF data to plot
 
     Returns
     -------
-    fig : plt.Figure
-        Output matplotlib figure
+    fig : matplotlib.figure.Figure
+        Output figure
     """
     fig, ax = plt.subplots()
 
diff --git a/src/gemdat/plots/_vibration.py b/src/gemdat/plots/_vibration.py
index 70f4b0fc..8eb1efa3 100644
--- a/src/gemdat/plots/_vibration.py
+++ b/src/gemdat/plots/_vibration.py
@@ -12,11 +12,11 @@ def frequency_vs_occurence(*, trajectory: Trajectory) -> plt.Figure:
     Parameters
     ----------
     trajectory : Trajectory
-        Input trajectory
+        Input trajectory, i.e. for the diffusing atom
 
     Returns
     -------
-    fig : plt.Figure
+    fig : matplotlib.figure.Figure
         Output figure
     """
     metrics = SimulationMetrics(trajectory)
@@ -68,11 +68,11 @@ def vibrational_amplitudes(*, trajectory: Trajectory) -> plt.Figure:
     Parameters
     ----------
     trajectory : Trajectory
-        Input trajectory
+        Input trajectory, i.e. for the diffusing atom
 
     Returns
     -------
-    fig : plt.Figure
+    fig : matplotlib.figure.Figure
         Output figure
     """
     metrics = SimulationMetrics(trajectory)
diff --git a/src/gemdat/rdf.py b/src/gemdat/rdf.py
index 03b33f8e..c7bbac09 100644
--- a/src/gemdat/rdf.py
+++ b/src/gemdat/rdf.py
@@ -73,7 +73,20 @@ def _get_symbol_indices(structure: Structure) -> dict[str, np.ndarray]:
 
 @dataclass
 class RDFData:
-    """Container for storing radial distribution data."""
+    """Data class for storing radial distribution data.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        1D array with x data (bins)
+    y : np.ndarray
+        1D array with y data (counts)
+    symbol : str
+        Distance to species with this symbol
+    state : str
+        State that the floating species is in, e.g.
+        the jump that it is making.
+    """
     x: np.ndarray
     y: np.ndarray
     symbol: str
@@ -81,10 +94,11 @@ class RDFData:
 
 
 def radial_distribution(
-        *,
-        sites: SitesData,
-        max_dist: float = 5.0,
-        resolution: float = 0.1) -> dict[str, dict[str, RDFData]]:
+    *,
+    sites: SitesData,
+    max_dist: float = 5.0,
+    resolution: float = 0.1,
+) -> dict[str, list[RDFData]]:
     """Calculate and sum RDFs for the floating species in the given sites data.
 
     Parameters
@@ -142,15 +156,17 @@ def radial_distribution(
                 rdfs[state_str, symbol] += np.bincount(rdf_state,
                                                        minlength=length)
 
-    ret: dict[str, dict[str, RDFData]] = defaultdict(dict)
+    ret: dict[str, list[RDFData]] = {}
 
     for (state, symbol), values in rdfs.items():
-        ret[state][symbol] = RDFData(
+        rdf_data = RDFData(
             x=bins,
             # Drop last element with distance > max_dist
             y=values[:-1],
             symbol=symbol,
             state=state,
         )
+        ret.setdefault(state, [])
+        ret[state].append(rdf_data)
 
     return ret
diff --git a/src/gemdat/simulation_metrics.py b/src/gemdat/simulation_metrics.py
index 6920971e..c192ffc9 100644
--- a/src/gemdat/simulation_metrics.py
+++ b/src/gemdat/simulation_metrics.py
@@ -1,3 +1,6 @@
+"""This module contains classes for calculating metrics and other properties
+from trajectories."""
+
 from __future__ import annotations
 
 import typing
@@ -14,8 +17,17 @@
 
 
 class SimulationMetrics:
+    """Class for calculating different metrics and properties from a molecular
+    dynamics simulation."""
 
     def __init__(self, trajectory: Trajectory):
+        """Initialize class.
+
+        Parameters
+        ----------
+        trajectory: Trajectory
+            Input trajectory
+        """
         self.trajectory = trajectory
 
     @weak_lru_cache()
@@ -23,13 +35,24 @@ def speed(self) -> np.ndarray:
         """Calculate speed.
 
         Corresponds to change in distance from the base position.
+
+        Returns
+        -------
+        speed : np.ndarray
+            Output array with speeds
         """
         distances = self.trajectory.distances_from_base_position()
         return np.diff(distances, prepend=0)
 
     @weak_lru_cache()
-    def particle_density(self) -> float:
-        """Return number of particles per m3."""
+    def particle_density(self) -> FloatWithUnit:
+        """Calculate number of particles per unit of volume from trajectory.
+
+        Returns
+        -------
+        particle_density : FloatWithUnit
+            Number of particles in $m^{-3}$
+        """
         lattice = self.trajectory.get_lattice()
         volume_ang = lattice.volume
         volume_m3 = volume_ang * angstrom**3
@@ -37,16 +60,22 @@ def particle_density(self) -> float:
         return FloatWithUnit(particle_density, 'm^-3')
 
     @weak_lru_cache()
-    def mol_per_liter(self) -> float:
-        """Return particle density as mol/liter."""
+    def mol_per_liter(self) -> FloatWithUnit:
+        """Calculate density.
+
+        Returns
+        -------
+        particle_density : FloatWithUnit
+            Particle density as $mol/l$.
+        """
         mol_per_liter = (self.particle_density() * 1e-3) / Avogadro
         return FloatWithUnit(mol_per_liter, 'mol l^-1')
 
     @weak_lru_cache()
-    def tracer_diffusivity(self, *, dimensions: int) -> float:
-        """Return tracer diffusivity in m2/s.
+    def tracer_diffusivity(self, *, dimensions: int) -> FloatWithUnit:
+        """Calculate tracer diffusivity.
 
-        Defined as: MSD/(2*dimensions*time)
+        Defined as: MSD / (2*dimensions*time)
 
         Parameters
         ----------
@@ -55,7 +84,8 @@ def tracer_diffusivity(self, *, dimensions: int) -> float:
 
         Returns
         -------
-        tracer_diffusivity : float
+        tracer_diffusivity : FloatWithUnit
+            Tracer diffusivity in $m^2/s$
         """
         distances = self.trajectory.distances_from_base_position()
         msd = np.mean(distances[:, -1]**2)  # Angstrom^2
@@ -66,7 +96,8 @@ def tracer_diffusivity(self, *, dimensions: int) -> float:
         return FloatWithUnit(tracer_diff, 'm^2 s^-1')
 
     @weak_lru_cache()
-    def tracer_conductivity(self, *, z_ion: int, dimensions: int) -> float:
+    def tracer_conductivity(self, *, z_ion: int,
+                            dimensions: int) -> FloatWithUnit:
         """Return tracer conductivity as S/m.
 
         Defined as: elementary_charge^2 * charge_ion^2 * diffusivity *
@@ -81,7 +112,8 @@ def tracer_conductivity(self, *, z_ion: int, dimensions: int) -> float:
 
         Returns
         -------
-        tracer_conductivity : float
+        tracer_conductivity : FloatWithUnit
+            Tracer conductivity in $S/m$
         """
         temperature = self.trajectory.metadata['temperature']
         tracer_diff = self.tracer_diffusivity(dimensions=dimensions)
@@ -91,14 +123,14 @@ def tracer_conductivity(self, *, z_ion: int, dimensions: int) -> float:
         return FloatWithUnit(tracer_conduc, 'S m^-1')
 
     @weak_lru_cache()
-    def attempt_frequency(self) -> tuple[float, float]:
+    def attempt_frequency(self) -> tuple[FloatWithUnit, FloatWithUnit]:
         """Return attempt frequency and standard deviation in Hz.
 
         Returns
         -------
-        attempt_freq : float
+        attempt_freq : FloatWithUnit
             Attempt frequency
-        attempt_freq_std : float
+        attempt_freq_std : FloatWithUnit
             Attempt frequency standard deviation
         """
         speed = self.speed()
@@ -114,12 +146,13 @@ def attempt_frequency(self) -> tuple[float, float]:
         return attempt_freq, attempt_freq_std
 
     @weak_lru_cache()
-    def vibration_amplitude(self) -> float:
-        """Return vibration amplitude in Angstrom.
+    def vibration_amplitude(self) -> FloatWithUnit:
+        """Calculate vibration amplitude.
 
         Returns
         -------
-        vibration_amp : float
+        vibration_amp : FloatWithUnit
+            Vibration amplitude in $Å$
         """
         amplitudes = self.amplitudes()
 
@@ -133,7 +166,7 @@ def vibration_amplitude(self) -> float:
 
     @weak_lru_cache()
     def amplitudes(self) -> np.ndarray:
-        """Return vibration amplitudes.
+        """Calculate vibration amplitudes.
 
         Returns
         -------
diff --git a/src/gemdat/sites.py b/src/gemdat/sites.py
index a9475c6d..6567ab7c 100644
--- a/src/gemdat/sites.py
+++ b/src/gemdat/sites.py
@@ -25,17 +25,19 @@
 
 class SitesData:
 
-    def __init__(self,
-                 *,
-                 structure: Structure,
-                 trajectory: Trajectory,
-                 floating_specie: str,
-                 n_parts: int = 10):
+    def __init__(
+        self,
+        *,
+        structure: Structure,
+        trajectory: Trajectory,
+        floating_specie: str,
+        n_parts: int = 10,
+    ):
         """Contain sites and jumps data.
 
         Parameters
         ----------
-        structure : Structure
+        structure : pymatgen.core.structure.Structure
             Input structure with known jump sites
         trajectory : Trajectory
             Input trajectory
@@ -64,43 +66,43 @@ def __init__(self,
     def set_split(self, n_parts: int):
         """Set number to split transitions into for statistics.
 
-        Sets `.parts`.
+        Sets [SitesData.parts][gemdat.sites.SitesData].
         """
         self.parts = self.transitions.split(n_parts,
                                             n_steps=len(self.trajectory))
 
     @property
-    def site_coords(self):
+    def site_coords(self) -> np.ndarray:
         """Return fractional coordinates for known sites."""
         return self.structure.frac_coords
 
     @property
-    def site_labels(self):
+    def site_labels(self) -> list[str]:
         """Return site labels."""
         return self.structure.labels
 
     @property
-    def n_sites(self):
+    def n_sites(self) -> int:
         """Return number of sites."""
         return len(self.structure)
 
     @property
-    def n_floating(self):
+    def n_floating(self) -> int:
         """Return number of floating species."""
         return len(self.diff_trajectory.species)
 
     @property
-    def n_jumps(self):
+    def n_jumps(self) -> int:
         """Return total number of jumps."""
         return len(self.transitions.events)
 
     @property
-    def n_solo_jumps(self):
+    def n_solo_jumps(self) -> int:
         """Return number of solo jumps."""
         return self.collective().n_solo_jumps
 
     @property
-    def solo_fraction(self):
+    def solo_fraction(self) -> float:
         """Fraction of solo jumps."""
         return self.n_solo_jumps / self.n_jumps
 
@@ -127,15 +129,21 @@ def site_pairs(self) -> list[tuple[str, str]]:
                      for label2 in labels})
 
     @property
-    def transitions_parts(self):
+    def transitions_parts(self) -> np.array:
+        """Return stacked array from
+        [part.matrix()][gemdat.transitions.Transitions.matrix]"""
         return np.stack([part.matrix() for part in self.parts])
 
     @property
-    def occupancy_parts(self):
+    def occupancy_parts(self) -> list[np.array]:
+        """Return [occupancy arrays][gemdat.transitions.Transitions.occupancy]
+        from parts."""
         return [part.occupancy() for part in self.parts]
 
     @property
-    def site_occupancy_parts(self):
+    def site_occupancy_parts(self) -> list[dict[str, float]]:
+        """Return [site occupancy][gemdat.sites.SitesData.site_occupancy] dicts
+        per part."""
         labels = self.site_labels
         n_steps = len(self.trajectory)
 
@@ -144,12 +152,14 @@ def site_occupancy_parts(self):
         return [
             _calculate_site_occupancy(occupancy=part.occupancy(),
                                       labels=labels,
-                                      n_steps=n_steps / self.n_parts)
+                                      n_steps=int(n_steps / self.n_parts))
             for part in parts
         ]
 
     @property
-    def atom_locations_parts(self):
+    def atom_locations_parts(self) -> 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
@@ -157,7 +167,8 @@ def atom_locations_parts(self):
         } for part in self.site_occupancy_parts]
 
     @property
-    def jumps_parts(self):
+    def jumps_parts(self) -> list[Counter]:
+        """Return [jump counters][gemdat.sites.SitesData.jumps] per part."""
         parts = self.parts
 
         labels = self.site_labels
@@ -265,8 +276,8 @@ def atom_locations(self):
         multiplier = self.n_sites / self.n_floating
         return {k: v * multiplier for k, v in self.site_occupancy().items()}
 
-    def jumps(self, ):
-        """Calculate number of jumpst between sites.
+    def jumps(self) -> Counter:
+        """Calculate number of jumps between sites.
 
         Returns
         -------
@@ -336,8 +347,12 @@ def collective(self, max_dist: float = 4.5) -> Collective:
         )
 
 
-def _calculate_site_occupancy(*, occupancy: dict[int, int], labels: list[str],
-                              n_steps: int) -> dict[str, float]:
+def _calculate_site_occupancy(
+    *,
+    occupancy: dict[int, int],
+    labels: list[str],
+    n_steps: int,
+) -> dict[str, float]:
     """Calculate percentage occupancy per unique site.
 
     Parameters
@@ -364,6 +379,5 @@ def _calculate_site_occupancy(*, occupancy: dict[int, int], labels: list[str],
 
     div = len(labels) * n_steps
     site_occupancies = {k: sum(v) / div for k, v in counts.items()}
-    # site_occupancies['total'] = sum(chain(*counts.values())) / (len(occupancy) * n_steps)
 
     return site_occupancies
diff --git a/src/gemdat/trajectory.py b/src/gemdat/trajectory.py
index 117c3321..aa0f6224 100644
--- a/src/gemdat/trajectory.py
+++ b/src/gemdat/trajectory.py
@@ -1,3 +1,8 @@
+"""This module contains class for dealing with trajectories from molecular
+dynamics simulations."""
+
+from __future__ import annotations
+
 import pickle
 from itertools import compress
 from pathlib import Path
@@ -33,8 +38,18 @@ def _lengths(vectors: np.ndarray, lattice: Lattice) -> np.ndarray:
 
 
 class Trajectory(PymatgenTrajectory):
+    """Trajectory of sites from a molecular dynamics simulation."""
 
     def __init__(self, *, metadata: dict | None = None, **kwargs):
+        """Initialize class.
+
+        Parameters
+        ----------
+        metadata : dict | None, optional
+            Optional dictionary with metadata
+        **kwargs
+            These are passed to [pymatgen.core.trajectory.Trajectory][]
+        """
         super().__init__(**kwargs)
         self.metadata = metadata if metadata else {}
 
@@ -49,7 +64,7 @@ def __getitem__(self, frames):
     def to_positions(self):
         """Pymatgen does not mod coords back to original unit cell.
 
-        https://github.com/GEMDAT-repos/GEMDAT/issues/103
+        See [GEMDAT#103](https://github.com/GEMDAT-repos/GEMDAT/issues/103)
         """
         super().to_positions()
         self.coords = np.mod(self.coords, 1)
@@ -90,7 +105,7 @@ def displacements(self) -> np.ndarray:
         return self.coords
 
     @classmethod
-    def from_cache(cls, cache: str | Path):
+    def from_cache(cls, cache: str | Path) -> Trajectory:
         """Load data from cache using pickle.
 
         Parameters
@@ -119,8 +134,8 @@ def from_vasprun(
         xml_file: str | Path,
         cache: Optional[str | Path] = None,
         **kwargs,
-    ):
-        """Load data from vasprun.xml.
+    ) -> Trajectory:
+        """Load data from a `vasprun.xml`.
 
         Parameters
         ----------
@@ -129,12 +144,12 @@ def from_vasprun(
         cache : Optional[Path], optional
             Path to cache data for vasprun.xml
         **kwargs : dict
-            Optional arguments passed to `vasp.Vasprun`
+            Optional arguments passed to [pymatgen.io.vasp.outputs.Vasprun][]
 
         Returns
         -------
-        data : Data
-            Dataclass with simulation data
+        trajectory : Trajectory
+            Output trajectory
         """
         kwargs.setdefault('parse_dos', False)
         kwargs.setdefault('parse_eigen', False)
@@ -178,7 +193,7 @@ def get_lattice(self, idx: int | None = None) -> Lattice:
 
         Returns
         -------
-        lattice : Lattice
+        lattice : pymatgen.core.lattice.Lattice
             Pymatgen Lattice object
         """
         if self.constant_lattice:
@@ -191,7 +206,8 @@ def get_lattice(self, idx: int | None = None) -> Lattice:
     def cumulative_displacements(self) -> np.ndarray:
         """Return cumulative displacement vectors from base positions.
 
-        This differs from `.displacements` in that it ignores the periodic boundary
+        This differs from [displacements()][gemdat.trajectory.Trajectory.displacements]
+        in that it ignores the periodic boundary
         conditions. Instead, it cumulatively tracks the lattice translation vector (jimage).
         """
         return np.cumsum(self.displacements, axis=0)
@@ -218,7 +234,7 @@ def drift(
         *,
         fixed_species: None | str | Collection[str] = None,
         floating_species: None | str | Collection[str] = None,
-    ):
+    ) -> np.ndarray:
         """Compute drift by averaging the displacement from the base positions
         per frame.
 
@@ -256,8 +272,9 @@ def apply_drift_correction(
         *,
         fixed_species: None | str | Collection[str] = None,
         floating_species: None | str | Collection[str] = None,
-    ):
-        """Apply drift correction to trajectory. For details see `Trajectory.drift`.
+    ) -> Trajectory:
+        """Apply drift correction to trajectory. For details see
+        [drift()][gemdat.trajectory.Trajectory.drift].
 
         If no species are specified, use all species to calculate drift.
 
@@ -286,7 +303,7 @@ def apply_drift_correction(
                               base_positions=self.base_positions,
                               time_step=self.time_step)
 
-    def filter(self, species: str | Collection[str]):
+    def filter(self, species: str | Collection[str]) -> Trajectory:
         """Return trajectory with coordinates for given species only.
 
         Parameters
diff --git a/src/gemdat/transitions.py b/src/gemdat/transitions.py
index 8b244970..425529c4 100644
--- a/src/gemdat/transitions.py
+++ b/src/gemdat/transitions.py
@@ -62,7 +62,7 @@ def from_trajectory(
         ----------
         trajectory : Trajectory
             Input trajectory
-        structure : Structure
+        structure : pymatgen.core.structure.Structure
             Input structure with known sites
         floating_specie : str
             Name of the floating specie to calculate transitions for
@@ -211,7 +211,7 @@ def _dist_close(trajectory: Trajectory, structure: Structure,
     ----------
     trajectory : Trajectory
         Input trajectory
-    structure : Structure
+    structure : pymatgen.core.structure.Structure
         Input structure
 
     Returns
@@ -268,7 +268,7 @@ def _calculate_atom_states(
 
     Parameters
     ----------
-    structure : Structure
+    structure : pymatgen.core.structure.Structure
         Input structure with pre-defined sites
     trajectory : Trajectory
         Input trajectory for floating atoms
diff --git a/src/gemdat/utils.py b/src/gemdat/utils.py
index 38d18edd..5cf6a1f9 100644
--- a/src/gemdat/utils.py
+++ b/src/gemdat/utils.py
@@ -1,3 +1,5 @@
+"""This module connects generally useful utilties."""
+
 import numpy as np
 from pymatgen.core import Lattice, Structure
 from scipy import signal
@@ -32,7 +34,7 @@ def ffill(arr: np.ndarray, fill_val: int = -1, axis=-1) -> np.ndarray:
 
 
 def bfill(arr: np.ndarray, fill_val: int = -1, axis=-1) -> np.ndarray:
-    """Backward fill.
+    """Backward fill values equal to `val` with upcoming values.
 
     See ffill for options.
     """
@@ -45,7 +47,7 @@ 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):
+def meanfreq(x: np.ndarray, fs: float = 1.0) -> np.ndarray:
     """Estimates the mean frequency in terms of the sample rate, fs.
 
     Vectorized version of https://stackoverflow.com/a/56487241
@@ -88,7 +90,7 @@ def is_lattice_similar(a: Lattice | Structure,
 
     Parameters
     ----------
-    a, b : Lattice | Structure
+    a, b : pymatgen.core.lattice.Lattice | pymatgen.core.structure.Structure
         Input lattices or structures
     length_tol : float, optional
         Length tolerance in Angstrom
diff --git a/src/gemdat/volume.py b/src/gemdat/volume.py
index 68873b28..0d260c3e 100644
--- a/src/gemdat/volume.py
+++ b/src/gemdat/volume.py
@@ -1,3 +1,5 @@
+"""This module contains functions related to dealing with volumetric data."""
+
 from __future__ import annotations
 
 from dataclasses import dataclass
@@ -21,13 +23,25 @@
 
 @dataclass
 class Volume:
+    """Container for volumetric data.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        Input volume as 3D numpy array
+    lattice : pymatgen.core.lattice.Lattice
+        Lattice parameters for the volume
+    resolution : optional[float]
+        The minimum resolution in Angstrom that the volume
+        was generated at.
+    """
     data: np.ndarray
     lattice: Lattice
-    # Minimum resolution in Angstrom
     resolution: float | None = None
 
     @property
-    def normalized_data(self):
+    def normalized_data(self) -> np.ndarray:
+        """Return normalized data."""
         return self.data / self.data.max()
 
     @property
@@ -36,20 +50,15 @@ def voxel_size(self) -> tuple[float, float, float]:
         return tuple(a / b for a, b in zip(self.lattice.lengths,
                                            self.data.shape))  # type: ignore
 
-    @property
-    def nx(self):
-        return self.data.shape[0]
-
-    @property
-    def ny(self):
-        return self.data.shape[1]
-
-    @property
-    def nz(self):
-        return self.data.shape[2]
-
     @classmethod
     def from_volumetric_data(cls, vol: VolumetricData):
+        """Create instance from VolumetricData.
+
+        Parameters
+        ----------
+        vol : pymatgen.io.common.VolumetricData
+            Input volumetric data
+        """
         return cls(data=vol.data['total'],
                    lattice=vol.structure.lattice,
                    resolution=None)
@@ -59,7 +68,11 @@ def trajectory_to_volume(
     trajectory: Trajectory,
     resolution: float = 0.2,
 ) -> np.ndarray:
-    """Calculate density volume from list of coordinates.
+    """Calculate density volume from a trajectory.
+
+    All coordinates are binned into voxels. The value of each
+    voxel represents the number of coodinates that are associated
+    with it.
 
     Parameters
     ----------
@@ -70,8 +83,8 @@ def trajectory_to_volume(
 
     Returns
     -------
-    vol : np.ndarray
-        3D numpy volume array
+    vol : Volume
+        Output volume
     """
     lattice = trajectory.get_lattice()
 
@@ -112,13 +125,13 @@ def trajectory_to_vasp_volume(trajectory: Trajectory,
                               structure: Optional[Structure] = None,
                               resolution: float = 0.2,
                               filename: str | None = None) -> VolumetricData:
-    """Calculate density volume as from list of coordinates.
+    """Calculate density volume from a trajectory.
 
     Parameters
     ----------
     trajectory : np.ndarray
         Input trajectory
-    structure : Optional[Structure]
+    structure : Optional[pymatgen.core.structure.Structure]
         structure to include in the vasp file, defaults to trajectory structure.
         Useful if you want to output the density for a select number of species,
         and show the host structure. Defaults to first structure in
@@ -130,14 +143,14 @@ def trajectory_to_vasp_volume(trajectory: Trajectory,
 
     Returns
     -------
-    vol : VolumetricData
+    vol : pymatgen.io.vasp.VolumetricData
         Output volumetric data object
     """
     vol = trajectory_to_volume(trajectory=trajectory, resolution=resolution)
 
     structure = structure if structure else trajectory.get_structure(0)
 
-    vasp_vol = VolumetricData(structure=structure, data={'total': vol})
+    vasp_vol = VolumetricData(structure=structure, data={'total': vol.data})
 
     if filename:
         vasp_vol.write_file(filename)
@@ -155,11 +168,14 @@ def volume_to_structure(
     vol_filename: str | None = None,
     **kwargs,
 ) -> Structure:
-    """Converts a volume array back to a structure using peak detection.
+    """Converts a volume back to a structure using peak detection.
+
+    Peaks are found using the Difference of Gaussian method available in
+    [scikit-image][skimage.feature.blob_dog].
 
     Parameters
     ----------
-    vol : Volume | VolumetricData
+    vol : Volume | pymatgen.io.common.VolumetricData
         Input volume data
     specie : str
         Specie to assign to the found sites
@@ -178,7 +194,7 @@ def volume_to_structure(
         If specified, write structure and volume in VASP format to this
         filename
     **kwargs
-        Additional keyword arguments are passed to skimage.feature.blob_dog
+        Additional keyword arguments are passed to [skimage.feature.blob_dog][]
     """
     if isinstance(vol, VolumetricData):
         vol = Volume.from_volumetric_data(vol)


Test patch:
diff --git a/tests/integration/plot_test.py b/tests/integration/plot_test.py
index 4ccba9e2..bf0ed008 100644
--- a/tests/integration/plot_test.py
+++ b/tests/integration/plot_test.py
@@ -32,34 +32,31 @@ def test_vibrational_amplitudes(vasp_traj):
 
 
 @image_comparison2(baseline_images=['jumps_vs_distance'])
-def test_jumps_vs_distance(vasp_traj, vasp_sites):
-    plots.jumps_vs_distance(trajectory=vasp_traj, sites=vasp_sites)
+def test_jumps_vs_distance(vasp_sites):
+    plots.jumps_vs_distance(sites=vasp_sites)
 
 
 @image_comparison2(baseline_images=['jumps_vs_time'])
-def test_jumps_vs_time(vasp_traj, vasp_sites):
-    plots.jumps_vs_time(trajectory=vasp_traj, sites=vasp_sites)
+def test_jumps_vs_time(vasp_sites):
+    plots.jumps_vs_time(sites=vasp_sites)
 
 
 @image_comparison2(baseline_images=['collective_jumps'])
-def test_collective_jumps(vasp_traj, vasp_sites):
-    plots.collective_jumps(trajectory=vasp_traj, sites=vasp_sites)
+def test_collective_jumps(vasp_sites):
+    plots.collective_jumps(sites=vasp_sites)
 
 
 @image_comparison2(baseline_images=['jumps_3d'])
-def test_jumps_3d(vasp_traj, vasp_sites):
-    plots.jumps_3d(trajectory=vasp_traj, sites=vasp_sites)
+def test_jumps_3d(vasp_sites):
+    plots.jumps_3d(sites=vasp_sites)
 
 
 @image_comparison2(baseline_images=['jumps_3d_animation'])
-def test_jumps_3d_animation(vasp_traj, vasp_sites):
-    plots.jumps_3d_animation(trajectory=vasp_traj,
-                             sites=vasp_sites,
-                             t_start=1000,
-                             t_stop=1001)
+def test_jumps_3d_animation(vasp_sites):
+    plots.jumps_3d_animation(sites=vasp_sites, t_start=1000, t_stop=1001)
 
 
 @image_comparison2(baseline_images=['rdf1', 'rdf2', 'rdf3'])
 def test_rdf(vasp_rdf_data):
     for rdfs in vasp_rdf_data.values():
-        plots.radial_distribution(rdfs.values())
+        plots.radial_distribution(rdfs)
diff --git a/tests/integration/rdf_test.py b/tests/integration/rdf_test.py
index 0dafd4a6..95410aee 100644
--- a/tests/integration/rdf_test.py
+++ b/tests/integration/rdf_test.py
@@ -10,6 +10,7 @@ def test_rdf(vasp_rdf_data):
 
     for state, rdfs in vasp_rdf_data.items():
         assert state in expected_states
-        assert set(rdfs.keys()) == expected_symbols
-        assert all(len(rdf.y) == 51 for rdf in rdfs.values())
-        assert all(len(rdf.x) == 51 for rdf in rdfs.values())
+        for rdf in rdfs:
+            assert rdf.symbol in expected_symbols
+            assert len(rdf.x) == 51
+            assert len(rdf.y) == 51