From a71973116e04fa83c2c8a67f9e3d9b2fbd32fd6b Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Thu, 8 Aug 2024 16:12:28 +0100
Subject: [PATCH 01/37] Adds component Mixin and Velocity osbervable

---
 janus_core/helpers/janus_types.py    |  32 +---
 janus_core/processing/correlator.py  |  35 +++-
 janus_core/processing/observables.py | 273 +++++++++++++++++++++++----
 tests/test_correlator.py             |  82 +++++++-
 4 files changed, 345 insertions(+), 77 deletions(-)

diff --git a/janus_core/helpers/janus_types.py b/janus_core/helpers/janus_types.py
index 708c7845..1c3db290 100644
--- a/janus_core/helpers/janus_types.py
+++ b/janus_core/helpers/janus_types.py
@@ -6,22 +6,15 @@
 from enum import Enum
 import logging
 from pathlib import Path, PurePath
-from typing import (
-    IO,
-    Literal,
-    Optional,
-    Protocol,
-    TypedDict,
-    TypeVar,
-    Union,
-    runtime_checkable,
-)
+from typing import IO, Literal, Optional, TypedDict, TypeVar, Union
 
 from ase import Atoms
 from ase.eos import EquationOfState
 import numpy as np
 from numpy.typing import NDArray
 
+from janus_core.helpers.observables import Observable
+
 # General
 
 T = TypeVar("T")
@@ -86,25 +79,6 @@ class PostProcessKwargs(TypedDict, total=False):
     vaf_output_file: PathLike | None
 
 
-@runtime_checkable
-class Observable(Protocol):
-    """Signature for correlation observable getter."""
-
-    def __call__(self, atoms: Atoms, *args, **kwargs) -> float:
-        """
-        Call the getter.
-
-        Parameters
-        ----------
-        atoms : Atoms
-            Atoms object to extract values from.
-        *args : tuple
-            Additional positional arguments passed to getter.
-        **kwargs : dict
-            Additional kwargs passed getter.
-        """
-
-
 class CorrelationKwargs(TypedDict, total=True):
     """Arguments for on-the-fly correlations <ab>."""
 
diff --git a/janus_core/processing/correlator.py b/janus_core/processing/correlator.py
index 39752a82..b3d8008b 100644
--- a/janus_core/processing/correlator.py
+++ b/janus_core/processing/correlator.py
@@ -7,7 +7,7 @@
 from ase import Atoms
 import numpy as np
 
-from janus_core.helpers.janus_types import Observable
+from janus_core.helpers.observables import Observable
 
 
 class Correlator:
@@ -238,7 +238,15 @@ def __init__(
             self._get_b = b
             self._b_args, self._b_kwargs = (), {}
 
-        self._correlator = Correlator(blocks=blocks, points=points, averaging=averaging)
+        self._correlators = []
+        for _ in zip(range(self._get_a.dimension), range(self._get_b.dimension)):
+            for _ in zip(
+                range(max(1, self._get_a.atom_count)),
+                range(max(1, self._get_b.atom_count)),
+            ):
+                self._correlators.append(
+                    Correlator(blocks=blocks, points=points, averaging=averaging)
+                )
         self._update_frequency = update_frequency
 
     @property
@@ -262,14 +270,17 @@ def update(self, atoms: Atoms) -> None:
         atoms : Atoms
             Atoms object to observe values from.
         """
-        self._correlator.update(
-            self._get_a(atoms, *self._a_args, **self._a_kwargs),
-            self._get_b(atoms, *self._b_args, **self._b_kwargs),
-        )
+        for i, values in enumerate(
+            zip(
+                self._get_a(atoms, *self._a_args, **self._a_kwargs),
+                self._get_b(atoms, *self._b_args, **self._b_kwargs),
+            )
+        ):
+            self._correlators[i].update(*values)
 
     def get(self) -> tuple[Iterable[float], Iterable[float]]:
         """
-        Get the correlation value and lags.
+        Get the correlation value and lags, averaging over atoms if applicable.
 
         Returns
         -------
@@ -278,7 +289,15 @@ def get(self) -> tuple[Iterable[float], Iterable[float]]:
         lags : Iterable[float]]
             The correlation lag times t'.
         """
-        return self._correlator.get()
+        if self._correlators:
+            avg_value, lags = self._correlators[0].get()
+            for cor in self._correlators[1:]:
+                value, _ = cor.get()
+                avg_value += value
+            return avg_value / max(
+                1, min(self._get_a.atom_count, self._get_b.atom_count)
+            ), lags
+        return [], []
 
     def __str__(self) -> str:
         """
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 32eee246..7fd4ba54 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -1,56 +1,210 @@
 """Module for built-in correlation observables."""
 
 from __future__ import annotations
+from typing import Optional
 
 from ase import Atoms, units
 
 
-class Stress:
+# pylint: disable=too-few-public-methods
+class Observable:
+    """
+    Observable data that may be correlated.
+
+    Parameters
+    ----------
+    dimension : int
+        The dimension of the observed data.
+    include_ideal_gas : bool
+            Calculate with the ideal gas contribution.
+    """
+
+    def __init__(self, component: str, *, include_ideal_gas: bool = True) -> None:
+        """
+        Initialise an observable with a given dimensionality.
+
+        Parameters
+        ----------
+        dimension : int
+            The dimension of the observed data.
+        include_ideal_gas : bool
+            Calculate with the ideal gas contribution.
+        """
+        self._dimension = dimension
+        self._getter = getter
+        self.atoms = None
+
+    def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
+        """
+        Call the user supplied getter if it exits.
+
+        Parameters
+        ----------
+        atoms : Atoms
+            Atoms object to extract values from.
+        *args : tuple
+            Additional positional arguments passed to getter.
+        **kwargs : dict
+            Additional kwargs passed getter.
+
+        Returns
+        -------
+        list[float]
+            The observed value, with dimensions atoms by self.dimension.
+
+        Raises
+        ------
+        ValueError
+            If user supplied getter is None.
+        """
+        if self._getter:
+            value = self._getter(atoms, *args, **kwargs)
+            if not isinstance(value, list):
+                return [value]
+            return value
+        raise ValueError("No user getter supplied")
+
+    @property
+    def dimension(self):
+        """
+        Dimension of the observable. Commensurate with self.__call__.
+
+        Returns
+        -------
+        int
+            Observables dimension.
+        """
+        return self._dimension
+
+    @property
+    def atom_count(self):
+        """
+        Atom count to average over.
+
+        Returns
+        -------
+        int
+            Atom count averaged over.
+        """
+        if self.atoms:
+            return len(self.atoms)
+        return 0
+
+
+class ComponentMixin:
+    """
+    Mixin to handle Observables with components.
+
+    Parameters
+    ----------
+    components : dict[str, int]
+        Symbolic components mapped to indices.
+    """
+
+    def __init__(self, components: dict[str, int]):
+        """
+        Initialise the mixin with components.
+
+        Parameters
+        ----------
+        components : dict[str, int]
+            Symbolic components mapped to indices.
+        """
+        self._components = components
+
+    @property
+    def allowed_components(self) -> dict[str, int]:
+        """
+        Allowed symbolic components with associated indices.
+
+        Returns
+        -------
+        Dict[str, int]
+            The allowed components and associated indices.
+        """
+        return self._components
+
+    @property
+    def _indices(self) -> list[int]:
+        """
+        Get indices associated with self._components.
+
+        Returns
+        -------
+        list[int]
+            The indices for each self._components.
+        """
+        return [self._components[c] for c in self.components]
+
+    def _set_components(self, components: list[str]):
+        """
+        Check if components are valid, if so set them.
+
+        Parameters
+        ----------
+        components : str
+            The component symbols to check.
+
+        Raises
+        ------
+        ValueError
+            If any component is invalid.
+        """
+        for component in components:
+            if component not in self.allowed_components:
+                component_names = list(self._components.keys())
+                raise ValueError(
+                    f"'{component}' invalid, must be '{', '.join(component_names)}'"
+                )
+        self.components = components
+
+
+# pylint: disable=too-few-public-methods
+class Stress(Observable, ComponentMixin):
     """
     Observable for stress components.
 
     Parameters
     ----------
-    component : str
-        Symbol for tensor components, xx, yy, etc.
+    components : list[str]
+        Symbols for correlated tensor components, xx, yy, etc.
     include_ideal_gas : bool
         Calculate with the ideal gas contribution.
     """
 
-    def __init__(self, component: str, *, include_ideal_gas: bool = True) -> None:
+    def __init__(self, components: list[str], *, include_ideal_gas: bool = True):
         """
-        Initialise the observables from a symbolic str component.
+        Initialise the observable from a symbolic str component.
 
         Parameters
         ----------
-        component : str
-            Symbol for tensor components, xx, yy, etc.
+        components : list[str]
+            Symbols for tensor components, xx, yy, etc.
         include_ideal_gas : bool
             Calculate with the ideal gas contribution.
         """
-        components = {
-            "xx": 0,
-            "yy": 1,
-            "zz": 2,
-            "yz": 3,
-            "zy": 3,
-            "xz": 4,
-            "zx": 4,
-            "xy": 5,
-            "yx": 5,
-        }
-        if component not in components:
-            raise ValueError(
-                f"'{component}' invalid, must be '{', '.join(list(components.keys()))}'"
-            )
-
-        self.component = component
-        self._index = components[self.component]
+        ComponentMixin.__init__(
+            self,
+            components={
+                "xx": 0,
+                "yy": 1,
+                "zz": 2,
+                "yz": 3,
+                "zy": 3,
+                "xz": 4,
+                "zx": 4,
+                "xy": 5,
+                "yx": 5,
+            },
+        )
+        self._set_components(components)
+
+        Observable.__init__(self, len(components))
         self.include_ideal_gas = include_ideal_gas
 
-    def __call__(self, atoms: Atoms, *args, **kwargs) -> float:
+    def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
         """
-        Get the stress component.
+        Get the stress components.
 
         Parameters
         ----------
@@ -63,12 +217,65 @@ def __call__(self, atoms: Atoms, *args, **kwargs) -> float:
 
         Returns
         -------
-        float
-            The stress component in GPa units.
+        list[float]
+            The stress components in GPa units.
         """
         return (
-            atoms.get_stress(include_ideal_gas=self.include_ideal_gas, voigt=True)[
-                self._index
-            ]
+            atoms.get_stress(include_ideal_gas=self.include_ideal_gas, voigt=True)
             / units.GPa
-        )
+        )[self._indices]
+
+
+StressDiagonal = Stress(["xx", "yy", "zz"])
+ShearStress = Stress(["xy", "yz", "zx"])
+
+
+# pylint: disable=too-few-public-methods
+class Velocity(Observable, ComponentMixin):
+    """
+    Observable for per atom velocity components.
+
+    Parameters
+    ----------
+    components : list[str]
+        Symbols for velocity components, x, y, z.
+    atoms : list[int]
+        List of atoms to observe velocities from.
+    """
+
+    def __init__(self, components: list[str], atoms: list[int]):
+        """
+        Initialise the observable from a symbolic str component and atom index.
+
+        Parameters
+        ----------
+        components : list[str]
+            Symbols for tensor components, x, y, and z.
+        atoms : list[int]
+            List of atoms to observe velocities from.
+        """
+        ComponentMixin.__init__(self, components={"x": 0, "y": 1, "z": 2})
+        self._set_components(components)
+
+        Observable.__init__(self, len(components))
+        self.atoms = atoms
+
+    def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
+        """
+        Get the velocity components for correlated atoms.
+
+        Parameters
+        ----------
+        atoms : Atoms
+            Atoms object to extract values from.
+        *args : tuple
+            Additional positional arguments passed to getter.
+        **kwargs : dict
+            Additional kwargs passed getter.
+
+        Returns
+        -------
+        list[float]
+            The velocity values.
+        """
+        return atoms.get_velocities()[self.atoms, :][:, self._indices].flatten()
diff --git a/tests/test_correlator.py b/tests/test_correlator.py
index 91b390e5..fae8c1f3 100644
--- a/tests/test_correlator.py
+++ b/tests/test_correlator.py
@@ -11,7 +11,7 @@
 import numpy as np
 from pytest import approx
 from typer.testing import CliRunner
-from yaml import Loader, load
+from yaml import Loader, load, safe_load
 
 from janus_core.calculations.md import NVE
 from janus_core.calculations.single_point import SinglePoint
@@ -66,6 +66,74 @@ def test_correlation():
     assert fft == approx(correlation, rel=1e-10)
 
 
+def test_vaf(tmp_path):
+    """Test the correlator against post-process."""
+    file_prefix = tmp_path / "Cl4Na4-nve-T300.0"
+    traj_path = tmp_path / "Cl4Na4-nve-T300.0-traj.extxyz"
+    cor_path = tmp_path / "Cl4Na4-nve-T300.0-cor.dat"
+
+    single_point = SinglePoint(
+        struct_path=DATA_PATH / "NaCl.cif",
+        arch="mace",
+        calc_kwargs={"model": MODEL_PATH},
+    )
+
+    na = []
+    cl = []
+    for i, atom in enumerate(single_point.struct):
+        if atom.symbol == "Na":
+            na.append(i)
+        else:
+            cl.append(i)
+
+    nve = NVE(
+        struct=single_point.struct,
+        temp=300.0,
+        steps=10,
+        seed=1,
+        traj_every=1,
+        stats_every=1,
+        file_prefix=file_prefix,
+        correlation_kwargs=[
+            {
+                "a": Velocity(["x", "y", "z"], na),
+                "b": Velocity(["x", "y", "z"], na),
+                "name": "vaf_Na",
+                "blocks": 1,
+                "points": 11,
+                "averaging": 1,
+                "update_frequency": 1,
+            },
+            {
+                "a": Velocity(["x", "y", "z"], cl),
+                "b": Velocity(["x", "y", "z"], cl),
+                "name": "vaf_Cl",
+                "blocks": 1,
+                "points": 11,
+                "averaging": 1,
+                "update_frequency": 1,
+            },
+        ],
+        write_kwargs={"invalidate_calc": False},
+    )
+
+    nve.run()
+
+    assert cor_path.exists()
+    assert traj_path.exists()
+
+    traj = read(traj_path, index=":")
+    vaf_post = post_process.compute_vaf(
+        traj, use_velocities=True, filter_atoms=(na, cl)
+    )
+    with open(cor_path) as cor:
+        vaf = safe_load(cor)
+    vaf_na = np.array(vaf["vaf_Na"]["value"])
+    vaf_cl = np.array(vaf["vaf_Cl"]["value"])
+    assert vaf_na == approx(vaf_post[0], rel=1e-5)
+    assert vaf_cl == approx(vaf_post[1], rel=1e-5)
+
+
 def test_md_correlations(tmp_path):
     """Test correlations as part of MD cycle."""
     file_prefix = tmp_path / "Cl4Na4-nve-T300.0"
@@ -78,10 +146,10 @@ def test_md_correlations(tmp_path):
         calc_kwargs={"model": MODEL_PATH},
     )
 
-    def user_observable_a(atoms: Atoms, kappa, **kwargs) -> float:
+    def user_observable_a(atoms: Atoms, kappa, *, gamma) -> float:
         """User specified getter for correlation."""
         return (
-            kwargs["gamma"]
+            gamma
             * kappa
             * atoms.get_stress(include_ideal_gas=True, voigt=True)[-1]
             / GPa
@@ -97,8 +165,8 @@ def user_observable_a(atoms: Atoms, kappa, **kwargs) -> float:
         file_prefix=file_prefix,
         correlation_kwargs=[
             {
-                "a": (user_observable_a, (2,), {"gamma": 2}),
-                "b": Stress("xy"),
+                "a": (Observable(1, getter=user_observable_a), (2,), {"gamma": 2}),
+                "b": Stress([("xy")]),
                 "name": "user_correlation",
                 "blocks": 1,
                 "points": 11,
@@ -106,8 +174,8 @@ def user_observable_a(atoms: Atoms, kappa, **kwargs) -> float:
                 "update_frequency": 1,
             },
             {
-                "a": Stress("xy"),
-                "b": Stress("xy"),
+                "a": Stress([("xy")]),
+                "b": Stress([("xy")]),
                 "name": "stress_xy_auto_cor",
                 "blocks": 1,
                 "points": 11,

From 33ec8f570f1286b6320a5781d83f8ec077b085c0 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Fri, 13 Sep 2024 14:45:22 +0100
Subject: [PATCH 02/37] Support SliceLike as atoms in Velocity

Also supports list[int], default is all atoms
---
 janus_core/calculations/md.py        |   4 +-
 janus_core/helpers/janus_types.py    | 122 ++++++++++++++++++++++++++-
 janus_core/processing/correlator.py  |  16 ++--
 janus_core/processing/observables.py |  24 ++++--
 tests/test_correlator.py             |   4 +-
 5 files changed, 152 insertions(+), 18 deletions(-)

diff --git a/janus_core/calculations/md.py b/janus_core/calculations/md.py
index 3b691795..379e0045 100644
--- a/janus_core/calculations/md.py
+++ b/janus_core/calculations/md.py
@@ -700,7 +700,9 @@ def _restart_file(self) -> str:
     def _parse_correlations(self) -> None:
         """Parse correlation kwargs into Correlations."""
         if self.correlation_kwargs:
-            self._correlations = [Correlation(**cor) for cor in self.correlation_kwargs]
+            self._correlations = [
+                Correlation(self.n_atoms, **cor) for cor in self.correlation_kwargs
+            ]
         else:
             self._correlations = ()
 
diff --git a/janus_core/helpers/janus_types.py b/janus_core/helpers/janus_types.py
index 1c3db290..63ac1cc8 100644
--- a/janus_core/helpers/janus_types.py
+++ b/janus_core/helpers/janus_types.py
@@ -13,8 +13,6 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from janus_core.helpers.observables import Observable
-
 # General
 
 T = TypeVar("T")
@@ -23,7 +21,6 @@
 PathLike = Union[str, Path]
 StartStopStep = tuple[Optional[int], Optional[int], int]
 SliceLike = Union[slice, range, int, StartStopStep]
-
 # ASE Arg types
 
 
@@ -155,3 +152,122 @@ class EoSResults(TypedDict, total=False):
     bulk_modulus: float
     v_0: float
     e_0: float
+
+
+# pylint: disable=too-few-public-methods
+class Observable:
+    """
+    Observable data that may be correlated.
+
+    Parameters
+    ----------
+    dimension : int
+        The dimension of the observed data.
+    getter : Optional[callable]
+        An optional callable to construct the Observable from.
+    """
+
+    def __init__(self, dimension: int = 1, *, getter: Optional[callable] = None):
+        """
+        Initialise an observable with a given dimensionality.
+
+        Parameters
+        ----------
+        dimension : int
+            The dimension of the observed data.
+        getter : Optional[callable]
+            An optional callable to construct the Observable from.
+        """
+        self._dimension = dimension
+        self._getter = getter
+        self.atoms = None
+
+    def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
+        """
+        Call the user supplied getter if it exits.
+
+        Parameters
+        ----------
+        atoms : Atoms
+            Atoms object to extract values from.
+        *args : tuple
+            Additional positional arguments passed to getter.
+        **kwargs : dict
+            Additional kwargs passed getter.
+
+        Returns
+        -------
+        list[float]
+            The observed value, with dimensions atoms by self.dimension.
+
+        Raises
+        ------
+        ValueError
+            If user supplied getter is None.
+        """
+        if self._getter:
+            value = self._getter(atoms, *args, **kwargs)
+            if not isinstance(value, list):
+                return [value]
+            return value
+        raise ValueError("No user getter supplied")
+
+    @property
+    def dimension(self):
+        """
+        Dimension of the observable. Commensurate with self.__call__.
+
+        Returns
+        -------
+        int
+            Observables dimension.
+        """
+        return self._dimension
+
+    def atom_count(self, n_atoms: int):
+        """
+        Atom count to average over.
+
+        Parameters
+        ----------
+        n_atoms : int
+            Total possible atoms.
+
+        Returns
+        -------
+        int
+            Atom count averaged over.
+        """
+        if self.atoms:
+            if isinstance(self.atoms, list):
+                return len(self.atoms)
+            if isinstance(self.atoms, int):
+                return 1
+
+            start = self.atoms.start
+            stop = self.atoms.stop
+            step = self.atoms.step
+            start = start if start is None else 0
+            stop = stop if stop is None else n_atoms
+            step = step if step is None else 1
+            return len(range(start, stop, step))
+        return 0
+
+
+class CorrelationKwargs(TypedDict, total=True):
+    """Arguments for on-the-fly correlations <ab>."""
+
+    #: observable a in <ab>, with optional args and kwargs
+    a: Union[Observable, tuple[Observable, tuple, dict]]
+    #: observable b in <ab>, with optional args and kwargs
+    b: Union[Observable, tuple[Observable, tuple, dict]]
+    #: name used for correlation in output
+    name: str
+    #: blocks used in multi-tau algorithm
+    blocks: int
+    #: points per block
+    points: int
+    #: averaging between blocks
+    averaging: int
+    #: frequency to update the correlation (steps)
+    update_frequency: int
diff --git a/janus_core/processing/correlator.py b/janus_core/processing/correlator.py
index b3d8008b..62b32e2e 100644
--- a/janus_core/processing/correlator.py
+++ b/janus_core/processing/correlator.py
@@ -179,6 +179,8 @@ class Correlation:
 
     Parameters
     ----------
+    n_atoms : int
+        Number of possible atoms to track.
     a : tuple[Observable, dict]
         Getter for a and kwargs.
     b : tuple[Observable, dict]
@@ -197,6 +199,7 @@ class Correlation:
 
     def __init__(
         self,
+        n_atoms: int,
         a: Observable | tuple[Observable, tuple, dict],
         b: Observable | tuple[Observable, tuple, dict],
         name: str,
@@ -210,6 +213,8 @@ def __init__(
 
         Parameters
         ----------
+        n_atoms : int
+            Number of possible atoms to track.
         a : tuple[Observable, tuple, dict]
             Getter for a and kwargs.
         b : tuple[Observable, tuple, dict]
@@ -238,11 +243,14 @@ def __init__(
             self._get_b = b
             self._b_args, self._b_kwargs = (), {}
 
+        self._a_atoms = self._get_a.atom_count(n_atoms)
+        self._b_atoms = self._get_b.atom_count(n_atoms)
+
         self._correlators = []
         for _ in zip(range(self._get_a.dimension), range(self._get_b.dimension)):
             for _ in zip(
-                range(max(1, self._get_a.atom_count)),
-                range(max(1, self._get_b.atom_count)),
+                range(max(1, self._a_atoms)),
+                range(max(1, self._b_atoms)),
             ):
                 self._correlators.append(
                     Correlator(blocks=blocks, points=points, averaging=averaging)
@@ -294,9 +302,7 @@ def get(self) -> tuple[Iterable[float], Iterable[float]]:
             for cor in self._correlators[1:]:
                 value, _ = cor.get()
                 avg_value += value
-            return avg_value / max(
-                1, min(self._get_a.atom_count, self._get_b.atom_count)
-            ), lags
+            return avg_value / max(1, min(self._a_atoms, self._b_atoms)), lags
         return [], []
 
     def __str__(self) -> str:
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 7fd4ba54..736bc26c 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -1,10 +1,12 @@
 """Module for built-in correlation observables."""
 
 from __future__ import annotations
-from typing import Optional
+
+from typing import Optional, Union
 
 from ase import Atoms, units
 
+from janus_core.helpers.janus_types import Observable, SliceLike
 
 # pylint: disable=too-few-public-methods
 class Observable:
@@ -90,7 +92,6 @@ def atom_count(self):
             return len(self.atoms)
         return 0
 
-
 class ComponentMixin:
     """
     Mixin to handle Observables with components.
@@ -239,11 +240,15 @@ class Velocity(Observable, ComponentMixin):
     ----------
     components : list[str]
         Symbols for velocity components, x, y, z.
-    atoms : list[int]
-        List of atoms to observe velocities from.
+    atoms : Optional[Union[list[int], SliceLike]]
+        List or slice of atoms to observe velocities from.
     """
 
-    def __init__(self, components: list[str], atoms: list[int]):
+    def __init__(
+        self,
+        components: list[str],
+        atoms: Optional[Union[list[int], SliceLike]] = None,
+    ):
         """
         Initialise the observable from a symbolic str component and atom index.
 
@@ -251,14 +256,17 @@ def __init__(self, components: list[str], atoms: list[int]):
         ----------
         components : list[str]
             Symbols for tensor components, x, y, and z.
-        atoms : list[int]
-            List of atoms to observe velocities from.
+        atoms : Union[list[int], SliceLike]
+            List or slice of atoms to observe velocities from.
         """
         ComponentMixin.__init__(self, components={"x": 0, "y": 1, "z": 2})
         self._set_components(components)
 
         Observable.__init__(self, len(components))
-        self.atoms = atoms
+        if atoms:
+            self.atoms = atoms
+        else:
+            atoms = slice(None, None, None)
 
     def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
         """
diff --git a/tests/test_correlator.py b/tests/test_correlator.py
index fae8c1f3..75e4242d 100644
--- a/tests/test_correlator.py
+++ b/tests/test_correlator.py
@@ -16,7 +16,9 @@
 from janus_core.calculations.md import NVE
 from janus_core.calculations.single_point import SinglePoint
 from janus_core.processing.correlator import Correlator
-from janus_core.processing.observables import Stress
+from janus_core.processing.observables import Stress, Velocity
+from janus_core.processing import post_process
+from janus_core.helpers.janus_types import Observable
 
 DATA_PATH = Path(__file__).parent / "data"
 MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model"

From 71afe8f947e3dea77f6d893bf6737173971f0a31 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Thu, 17 Oct 2024 10:28:28 +0100
Subject: [PATCH 03/37] Move Observable into observables.py

---
 janus_core/calculations/md.py        |   3 +-
 janus_core/helpers/janus_types.py    | 102 +------------------------
 janus_core/processing/correlator.py  |   9 ++-
 janus_core/processing/observables.py | 107 ++++++++++++++++++++++++++-
 4 files changed, 113 insertions(+), 108 deletions(-)

diff --git a/janus_core/calculations/md.py b/janus_core/calculations/md.py
index 379e0045..bfc5573d 100644
--- a/janus_core/calculations/md.py
+++ b/janus_core/calculations/md.py
@@ -701,7 +701,8 @@ def _parse_correlations(self) -> None:
         """Parse correlation kwargs into Correlations."""
         if self.correlation_kwargs:
             self._correlations = [
-                Correlation(self.n_atoms, **cor) for cor in self.correlation_kwargs
+                Correlation(n_atoms=self.n_atoms, **cor)
+                for cor in self.correlation_kwargs
             ]
         else:
             self._correlations = ()
diff --git a/janus_core/helpers/janus_types.py b/janus_core/helpers/janus_types.py
index 63ac1cc8..d000ebbb 100644
--- a/janus_core/helpers/janus_types.py
+++ b/janus_core/helpers/janus_types.py
@@ -13,6 +13,8 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from janus_core.helpers.observables import Observable
+
 # General
 
 T = TypeVar("T")
@@ -154,106 +156,6 @@ class EoSResults(TypedDict, total=False):
     e_0: float
 
 
-# pylint: disable=too-few-public-methods
-class Observable:
-    """
-    Observable data that may be correlated.
-
-    Parameters
-    ----------
-    dimension : int
-        The dimension of the observed data.
-    getter : Optional[callable]
-        An optional callable to construct the Observable from.
-    """
-
-    def __init__(self, dimension: int = 1, *, getter: Optional[callable] = None):
-        """
-        Initialise an observable with a given dimensionality.
-
-        Parameters
-        ----------
-        dimension : int
-            The dimension of the observed data.
-        getter : Optional[callable]
-            An optional callable to construct the Observable from.
-        """
-        self._dimension = dimension
-        self._getter = getter
-        self.atoms = None
-
-    def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
-        """
-        Call the user supplied getter if it exits.
-
-        Parameters
-        ----------
-        atoms : Atoms
-            Atoms object to extract values from.
-        *args : tuple
-            Additional positional arguments passed to getter.
-        **kwargs : dict
-            Additional kwargs passed getter.
-
-        Returns
-        -------
-        list[float]
-            The observed value, with dimensions atoms by self.dimension.
-
-        Raises
-        ------
-        ValueError
-            If user supplied getter is None.
-        """
-        if self._getter:
-            value = self._getter(atoms, *args, **kwargs)
-            if not isinstance(value, list):
-                return [value]
-            return value
-        raise ValueError("No user getter supplied")
-
-    @property
-    def dimension(self):
-        """
-        Dimension of the observable. Commensurate with self.__call__.
-
-        Returns
-        -------
-        int
-            Observables dimension.
-        """
-        return self._dimension
-
-    def atom_count(self, n_atoms: int):
-        """
-        Atom count to average over.
-
-        Parameters
-        ----------
-        n_atoms : int
-            Total possible atoms.
-
-        Returns
-        -------
-        int
-            Atom count averaged over.
-        """
-        if self.atoms:
-            if isinstance(self.atoms, list):
-                return len(self.atoms)
-            if isinstance(self.atoms, int):
-                return 1
-
-            start = self.atoms.start
-            stop = self.atoms.stop
-            step = self.atoms.step
-            start = start if start is None else 0
-            stop = stop if stop is None else n_atoms
-            step = step if step is None else 1
-            return len(range(start, stop, step))
-        return 0
-
-
 class CorrelationKwargs(TypedDict, total=True):
     """Arguments for on-the-fly correlations <ab>."""
 
diff --git a/janus_core/processing/correlator.py b/janus_core/processing/correlator.py
index 62b32e2e..2921b9e1 100644
--- a/janus_core/processing/correlator.py
+++ b/janus_core/processing/correlator.py
@@ -181,9 +181,9 @@ class Correlation:
     ----------
     n_atoms : int
         Number of possible atoms to track.
-    a : tuple[Observable, dict]
+    a : Union[Observable, tuple[Observable, tuple, dict]]
         Getter for a and kwargs.
-    b : tuple[Observable, dict]
+    b : Union[Observable, tuple[Observable, tuple, dict]]
         Getter for b and kwargs.
     name : str
         Name of correlation.
@@ -199,6 +199,7 @@ class Correlation:
 
     def __init__(
         self,
+        *,
         n_atoms: int,
         a: Observable | tuple[Observable, tuple, dict],
         b: Observable | tuple[Observable, tuple, dict],
@@ -215,9 +216,9 @@ def __init__(
         ----------
         n_atoms : int
             Number of possible atoms to track.
-        a : tuple[Observable, tuple, dict]
+        a : Union[Observable, tuple[Observable, tuple, dict]]
             Getter for a and kwargs.
-        b : tuple[Observable, tuple, dict]
+        b : Union[Observable, tuple[Observable, tuple, dict]]
             Getter for b and kwargs.
         name : str
             Name of correlation.
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 736bc26c..bd1855c6 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -2,11 +2,112 @@
 
 from __future__ import annotations
 
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from ase import Atoms, units
 
-from janus_core.helpers.janus_types import Observable, SliceLike
+if TYPE_CHECKING:
+    from janus_core.helpers.janus_types import SliceLike
+
+
+# pylint: disable=too-few-public-methods
+class Observable:
+    """
+    Observable data that may be correlated.
+
+    Parameters
+    ----------
+    dimension : int
+        The dimension of the observed data.
+    getter : Optional[callable]
+        An optional callable to construct the Observable from.
+    """
+
+    def __init__(self, dimension: int = 1, *, getter: Optional[callable] = None):
+        """
+        Initialise an observable with a given dimensionality.
+
+        Parameters
+        ----------
+        dimension : int
+            The dimension of the observed data.
+        getter : Optional[callable]
+            An optional callable to construct the Observable from.
+        """
+        self._dimension = dimension
+        self._getter = getter
+        self.atoms = None
+
+    def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
+        """
+        Call the user supplied getter if it exits.
+
+        Parameters
+        ----------
+        atoms : Atoms
+            Atoms object to extract values from.
+        *args : tuple
+            Additional positional arguments passed to getter.
+        **kwargs : dict
+            Additional kwargs passed getter.
+
+        Returns
+        -------
+        list[float]
+            The observed value, with dimensions atoms by self.dimension.
+
+        Raises
+        ------
+        ValueError
+            If user supplied getter is None.
+        """
+        if self._getter:
+            value = self._getter(atoms, *args, **kwargs)
+            if not isinstance(value, list):
+                return [value]
+            return value
+        raise ValueError("No user getter supplied")
+
+    @property
+    def dimension(self):
+        """
+        Dimension of the observable. Commensurate with self.__call__.
+
+        Returns
+        -------
+        int
+            Observables dimension.
+        """
+        return self._dimension
+
+    def atom_count(self, n_atoms: int):
+        """
+        Atom count to average over.
+
+        Parameters
+        ----------
+        n_atoms : int
+            Total possible atoms.
+
+        Returns
+        -------
+        int
+            Atom count averaged over.
+        """
+        if self.atoms:
+            if isinstance(self.atoms, list):
+                return len(self.atoms)
+            if isinstance(self.atoms, int):
+                return 1
+
+            start = self.atoms.start
+            stop = self.atoms.stop
+            step = self.atoms.step
+            start = start if start is None else 0
+            stop = stop if stop is None else n_atoms
+            step = step if step is None else 1
+            return len(range(start, stop, step))
+        return 0
 
 # pylint: disable=too-few-public-methods
 class Observable:
@@ -247,7 +348,7 @@ class Velocity(Observable, ComponentMixin):
     def __init__(
         self,
         components: list[str],
-        atoms: Optional[Union[list[int], SliceLike]] = None,
+        atoms: Optional[Union[list[int], "SliceLike"]] = None,
     ):
         """
         Initialise the observable from a symbolic str component and atom index.

From 626fbe2f2e950b3e8b8a2e2dbf5f7ad25b8fbd4b Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Tue, 22 Oct 2024 11:07:33 +0100
Subject: [PATCH 04/37] Remove getter

---
 janus_core/processing/observables.py | 20 ++----------------
 tests/test_correlator.py             | 31 +---------------------------
 2 files changed, 3 insertions(+), 48 deletions(-)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index bd1855c6..0a331d36 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -19,11 +19,9 @@ class Observable:
     ----------
     dimension : int
         The dimension of the observed data.
-    getter : Optional[callable]
-        An optional callable to construct the Observable from.
     """
 
-    def __init__(self, dimension: int = 1, *, getter: Optional[callable] = None):
+    def __init__(self, dimension: int = 1):
         """
         Initialise an observable with a given dimensionality.
 
@@ -31,16 +29,13 @@ def __init__(self, dimension: int = 1, *, getter: Optional[callable] = None):
         ----------
         dimension : int
             The dimension of the observed data.
-        getter : Optional[callable]
-            An optional callable to construct the Observable from.
         """
         self._dimension = dimension
-        self._getter = getter
         self.atoms = None
 
     def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
         """
-        Call the user supplied getter if it exits.
+        Signature for returning observed value from atoms.
 
         Parameters
         ----------
@@ -55,18 +50,7 @@ def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
         -------
         list[float]
             The observed value, with dimensions atoms by self.dimension.
-
-        Raises
-        ------
-        ValueError
-            If user supplied getter is None.
         """
-        if self._getter:
-            value = self._getter(atoms, *args, **kwargs)
-            if not isinstance(value, list):
-                return [value]
-            return value
-        raise ValueError("No user getter supplied")
 
     @property
     def dimension(self):
diff --git a/tests/test_correlator.py b/tests/test_correlator.py
index 75e4242d..a6c5f03e 100644
--- a/tests/test_correlator.py
+++ b/tests/test_correlator.py
@@ -5,7 +5,6 @@
 from collections.abc import Iterable
 from pathlib import Path
 
-from ase import Atoms
 from ase.io import read
 from ase.units import GPa
 import numpy as np
@@ -18,7 +17,6 @@
 from janus_core.processing.correlator import Correlator
 from janus_core.processing.observables import Stress, Velocity
 from janus_core.processing import post_process
-from janus_core.helpers.janus_types import Observable
 
 DATA_PATH = Path(__file__).parent / "data"
 MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model"
@@ -148,15 +146,6 @@ def test_md_correlations(tmp_path):
         calc_kwargs={"model": MODEL_PATH},
     )
 
-    def user_observable_a(atoms: Atoms, kappa, *, gamma) -> float:
-        """User specified getter for correlation."""
-        return (
-            gamma
-            * kappa
-            * atoms.get_stress(include_ideal_gas=True, voigt=True)[-1]
-            / GPa
-        )
-
     nve = NVE(
         struct=single_point.struct,
         temp=300.0,
@@ -166,15 +155,6 @@ def user_observable_a(atoms: Atoms, kappa, *, gamma) -> float:
         stats_every=1,
         file_prefix=file_prefix,
         correlation_kwargs=[
-            {
-                "a": (Observable(1, getter=user_observable_a), (2,), {"gamma": 2}),
-                "b": Stress([("xy")]),
-                "name": "user_correlation",
-                "blocks": 1,
-                "points": 11,
-                "averaging": 1,
-                "update_frequency": 1,
-            },
             {
                 "a": Stress([("xy")]),
                 "b": Stress([("xy")]),
@@ -197,8 +177,7 @@ def user_observable_a(atoms: Atoms, kappa, *, gamma) -> float:
     assert cor_path.exists()
     with open(cor_path, encoding="utf8") as in_file:
         cor = load(in_file, Loader=Loader)
-    assert len(cor) == 2
-    assert "user_correlation" in cor
+    assert len(cor) == 1
     assert "stress_xy_auto_cor" in cor
 
     stress_cor = cor["stress_xy_auto_cor"]
@@ -208,11 +187,3 @@ def user_observable_a(atoms: Atoms, kappa, *, gamma) -> float:
     direct = correlate(pxy, pxy, fft=False)
     # input data differs due to i/o, error is expected 1e-5
     assert direct == approx(value, rel=1e-5)
-
-    user_cor = cor["user_correlation"]
-    value, lags = user_cor["value"], stress_cor["lags"]
-    assert len(value) == len(lags) == 11
-
-    direct = correlate([v * 4.0 for v in pxy], pxy, fft=False)
-    # input data differs due to i/o, error is expected 1e-5
-    assert direct == approx(value, rel=1e-5)

From d706c924745c5ef31ca788282f3d55d183d84418 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Tue, 22 Oct 2024 12:08:36 +0100
Subject: [PATCH 05/37] USe annotations

---
 janus_core/processing/observables.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 0a331d36..2a516c6c 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 from ase import Atoms, units
 
@@ -332,7 +332,7 @@ class Velocity(Observable, ComponentMixin):
     def __init__(
         self,
         components: list[str],
-        atoms: Optional[Union[list[int], "SliceLike"]] = None,
+        atoms: list[int] | SliceLike | None = None,
     ):
         """
         Initialise the observable from a symbolic str component and atom index.

From 772f198b1d7b06b9e986dbee898dc771fcc5654f Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Tue, 22 Oct 2024 16:20:51 +0100
Subject: [PATCH 06/37] Create SliceLike utils

---
 janus_core/helpers/utils.py           | 52 +++++++++++++++++++++++++--
 janus_core/processing/observables.py  | 12 ++-----
 janus_core/processing/post_process.py | 31 ++--------------
 tests/test_utils.py                   | 41 +++++++++++++++++++++
 4 files changed, 96 insertions(+), 40 deletions(-)

diff --git a/janus_core/helpers/utils.py b/janus_core/helpers/utils.py
index 3153fac4..f4d28101 100644
--- a/janus_core/helpers/utils.py
+++ b/janus_core/helpers/utils.py
@@ -18,8 +18,7 @@
 )
 from rich.style import Style
 
-from janus_core.helpers.janus_types import MaybeSequence, PathLike
-
+from janus_core.helpers.janus_types import MaybeSequence, PathLike, SliceLike, StartStopStep
 
 class FileNameMixin(ABC):  # noqa: B024 (abstract-base-class-without-abstract-method)
     """
@@ -409,3 +408,52 @@ def track_progress(sequence: Sequence | Iterable, description: str) -> Iterable:
 
     with progress:
         yield from progress.track(sequence, description=description)
+
+def slicelike_to_startstopstep(index: SliceLike) -> StartStopStep:
+    """
+    Standarize `SliceLike`s into tuple of `start`, `stop`, `step`.
+
+    Parameters
+    ----------
+    index : SliceLike
+        `SliceLike` to standardize.
+
+    Returns
+    -------
+    StartStopStep
+        Standardized `SliceLike` as `start`, `stop`, `step` triplet.
+    """
+    if isinstance(index, int):
+        if index == -1:
+            return (index, None, 1)
+        return (index, index + 1, 1)
+
+    if isinstance(index, (slice, range)):
+        return (index.start, index.stop, index.step)
+
+    return index
+
+
+def slicelike_len_for(slc: SliceLike, sliceable_length: int) -> int:
+    """
+    Calculate the length of a SliceLike applied to a sliceable of a given length.
+
+    Parameters
+    ----------
+    slc : SliceLike
+        The applied SliceLike.
+    sliceable_length : int
+        The length of the sliceable object.
+
+    Returns
+    -------
+    int
+        Length of the result of applying slc.
+    """
+    start, stop, step = slicelike_to_startstopstep(slc)
+    if stop is None:
+        stop = sliceable_length
+    # start = start if start is None else 0
+    # stop = stop if stop is None else sliceable_length
+    # step = step if step is None else 1
+    return len(range(start, stop, step))
\ No newline at end of file
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 2a516c6c..7b8f1213 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -8,6 +8,7 @@
 
 if TYPE_CHECKING:
     from janus_core.helpers.janus_types import SliceLike
+    from janus_core.helpers.utils import slicelike_len_for
 
 
 # pylint: disable=too-few-public-methods
@@ -81,16 +82,7 @@ def atom_count(self, n_atoms: int):
         if self.atoms:
             if isinstance(self.atoms, list):
                 return len(self.atoms)
-            if isinstance(self.atoms, int):
-                return 1
-
-            start = self.atoms.start
-            stop = self.atoms.stop
-            step = self.atoms.step
-            start = start if start is None else 0
-            stop = stop if stop is None else n_atoms
-            step = step if step is None else 1
-            return len(range(start, stop, step))
+            return slicelike_len_for(self.n_atoms)
         return 0
 
 # pylint: disable=too-few-public-methods
diff --git a/janus_core/processing/post_process.py b/janus_core/processing/post_process.py
index 044b8d02..2c9a66ff 100644
--- a/janus_core/processing/post_process.py
+++ b/janus_core/processing/post_process.py
@@ -15,33 +15,8 @@
     MaybeSequence,
     PathLike,
     SliceLike,
-    StartStopStep,
 )
-
-
-def _process_index(index: SliceLike) -> StartStopStep:
-    """
-    Standarize `SliceLike`s into tuple of `start`, `stop`, `step`.
-
-    Parameters
-    ----------
-    index : SliceLike
-        `SliceLike` to standardize.
-
-    Returns
-    -------
-    StartStopStep
-        Standardized `SliceLike` as `start`, `stop`, `step` triplet.
-    """
-    if isinstance(index, int):
-        if index == -1:
-            return (index, None, 1)
-        return (index, index + 1, 1)
-
-    if isinstance(index, (slice, range)):
-        return (index.start, index.stop, index.step)
-
-    return index
+from janus_core.helpers.utils import slicelike_to_startstopstep
 
 
 def compute_rdf(
@@ -94,7 +69,7 @@ def compute_rdf(
         If `by_elements` is true returns a `dict` of RDF by element pairs.
         Otherwise returns RDF of total system filtered by elements.
     """
-    index = _process_index(index)
+    index = slicelike_to_startstopstep(index)
 
     if not isinstance(data, Sequence):
         data = [data]
@@ -261,7 +236,7 @@ def compute_vaf(
             )
 
     # Extract requested data
-    index = _process_index(index)
+    index = slicelike_to_startstopstep(index)
     data = data[slice(*index)]
 
     if use_velocities:
diff --git a/tests/test_utils.py b/tests/test_utils.py
index c000b6af..03cd2915 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -9,9 +9,16 @@
 import pytest
 
 from janus_core.cli.utils import dict_paths_to_strs, dict_remove_hyphens
+from janus_core.helpers.janus_types import SliceLike, StartStopStep
 from janus_core.helpers.mlip_calculators import choose_calculator
 from janus_core.helpers.struct_io import output_structs
 from janus_core.helpers.utils import none_to_dict
+from janus_core.helpers.utils import (
+    none_to_dict,
+    output_structs,
+    slicelike_len_for,
+    slicelike_to_startstopstep,
+)
 
 DATA_PATH = Path(__file__).parent / "data/NaCl.cif"
 MODEL_PATH = Path(__file__).parent / "models/mace_mp_small.model"
@@ -157,3 +164,37 @@ def test_none_to_dict(dicts_in):
     assert dicts[2] == dicts_in[2]
     assert dicts[3] == dicts_in[3]
     assert dicts[4] == {}
+
+
+@pytest.mark.parametrize(
+    "slc, expected",
+    [
+        ((1, 2, 3), (1, 2, 3)),
+        (1, (1, 2, 1)),
+        (range(1, 2, 3), (1, 2, 3)),
+        (slice(1, 2, 3), (1, 2, 3)),
+        (-1, (-1, None, 1)),
+        (range(10), (0, 10, 1)),
+        (slice(0, None, 1), (0, None, 1)),
+    ],
+)
+def test_slicelike_to_startstopstep(slc: SliceLike, expected: StartStopStep):
+    """Test converting SliceLike to StartStopStep."""
+    assert slicelike_to_startstopstep(slc) == expected
+
+
+@pytest.mark.parametrize(
+    "slc_len, expected",
+    [
+        (((1, 2, 3), 3), 1),
+        ((1, 1), 1),
+        ((range(1, 2, 3), 3), 1),
+        ((slice(1, 2, 3), 3), 1),
+        ((-1, 1), 2),
+        ((range(10), 10), 10),
+        ((slice(0, None, 2), 10), 5),
+    ],
+)
+def test_slicelike_len_for(slc_len: tuple[SliceLike, int], expected: int):
+    """Test converting SliceLike to StartStopStep."""
+    assert slicelike_len_for(*slc_len) == expected

From d0367cd9e51bad76539ab16d6c13b2993c4c25ba Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Wed, 23 Oct 2024 15:01:54 +0100
Subject: [PATCH 07/37] Clarify values

Use kwargs only in inits

atoms -> atoms_slice

replace atom_count method with value_count method

Stress now can slice atoms
---
 janus_core/processing/correlator.py  |  28 +++--
 janus_core/processing/observables.py | 161 ++++++++++-----------------
 tests/test_correlator.py             |  16 +--
 3 files changed, 77 insertions(+), 128 deletions(-)

diff --git a/janus_core/processing/correlator.py b/janus_core/processing/correlator.py
index 2921b9e1..d9e9eb7d 100644
--- a/janus_core/processing/correlator.py
+++ b/janus_core/processing/correlator.py
@@ -244,18 +244,18 @@ def __init__(
             self._get_b = b
             self._b_args, self._b_kwargs = (), {}
 
-        self._a_atoms = self._get_a.atom_count(n_atoms)
-        self._b_atoms = self._get_b.atom_count(n_atoms)
+        a_values = self._get_a.value_count(n_atoms)
+        b_values = self._get_b.value_count(n_atoms)
+
+        if a_values != b_values:
+            raise ValueError("Observables have inconsistent sizes")
+        self._values = a_values
 
         self._correlators = []
-        for _ in zip(range(self._get_a.dimension), range(self._get_b.dimension)):
-            for _ in zip(
-                range(max(1, self._a_atoms)),
-                range(max(1, self._b_atoms)),
-            ):
-                self._correlators.append(
-                    Correlator(blocks=blocks, points=points, averaging=averaging)
-                )
+        for _ in range(self._values):
+            self._correlators.append(
+                Correlator(blocks=blocks, points=points, averaging=averaging)
+            )
         self._update_frequency = update_frequency
 
     @property
@@ -299,11 +299,9 @@ def get(self) -> tuple[Iterable[float], Iterable[float]]:
             The correlation lag times t'.
         """
         if self._correlators:
-            avg_value, lags = self._correlators[0].get()
-            for cor in self._correlators[1:]:
-                value, _ = cor.get()
-                avg_value += value
-            return avg_value / max(1, min(self._a_atoms, self._b_atoms)), lags
+            _, lags = self._correlators[0].get()
+            avg_value = sum([cor.get()[0] for cor in self._correlators]) / self._values
+            return avg_value, lags
         return [], []
 
     def __str__(self) -> str:
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 7b8f1213..d17b8b2b 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -32,7 +32,6 @@ def __init__(self, dimension: int = 1):
             The dimension of the observed data.
         """
         self._dimension = dimension
-        self.atoms = None
 
     def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
         """
@@ -53,95 +52,21 @@ def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
             The observed value, with dimensions atoms by self.dimension.
         """
 
-    @property
-    def dimension(self):
-        """
-        Dimension of the observable. Commensurate with self.__call__.
-
-        Returns
-        -------
-        int
-            Observables dimension.
-        """
-        return self._dimension
-
-    def atom_count(self, n_atoms: int):
+    def value_count(self, n_atoms: int | None = None) -> int:
         """
-        Atom count to average over.
+        Count of values returned by __call__.
 
         Parameters
         ----------
-        n_atoms : int
-            Total possible atoms.
+        n_atoms : int | None
+            Atom count to expand atoms_slice.
 
         Returns
         -------
         int
-            Atom count averaged over.
+            The number of values returned by __call__.
         """
-        if self.atoms:
-            if isinstance(self.atoms, list):
-                return len(self.atoms)
-            return slicelike_len_for(self.n_atoms)
-        return 0
-
-# pylint: disable=too-few-public-methods
-class Observable:
-    """
-    Observable data that may be correlated.
-
-    Parameters
-    ----------
-    dimension : int
-        The dimension of the observed data.
-    include_ideal_gas : bool
-            Calculate with the ideal gas contribution.
-    """
-
-    def __init__(self, component: str, *, include_ideal_gas: bool = True) -> None:
-        """
-        Initialise an observable with a given dimensionality.
-
-        Parameters
-        ----------
-        dimension : int
-            The dimension of the observed data.
-        include_ideal_gas : bool
-            Calculate with the ideal gas contribution.
-        """
-        self._dimension = dimension
-        self._getter = getter
-        self.atoms = None
-
-    def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
-        """
-        Call the user supplied getter if it exits.
-
-        Parameters
-        ----------
-        atoms : Atoms
-            Atoms object to extract values from.
-        *args : tuple
-            Additional positional arguments passed to getter.
-        **kwargs : dict
-            Additional kwargs passed getter.
-
-        Returns
-        -------
-        list[float]
-            The observed value, with dimensions atoms by self.dimension.
-
-        Raises
-        ------
-        ValueError
-            If user supplied getter is None.
-        """
-        if self._getter:
-            value = self._getter(atoms, *args, **kwargs)
-            if not isinstance(value, list):
-                return [value]
-            return value
-        raise ValueError("No user getter supplied")
+        return self.dimension
 
     @property
     def dimension(self):
@@ -155,19 +80,6 @@ def dimension(self):
         """
         return self._dimension
 
-    @property
-    def atom_count(self):
-        """
-        Atom count to average over.
-
-        Returns
-        -------
-        int
-            Atom count averaged over.
-        """
-        if self.atoms:
-            return len(self.atoms)
-        return 0
 
 class ComponentMixin:
     """
@@ -246,11 +158,19 @@ class Stress(Observable, ComponentMixin):
     ----------
     components : list[str]
         Symbols for correlated tensor components, xx, yy, etc.
+    atoms_slice : list[int] | SliceLike | None = None
+        List or slice of atoms to observe velocities from.
     include_ideal_gas : bool
         Calculate with the ideal gas contribution.
     """
 
-    def __init__(self, components: list[str], *, include_ideal_gas: bool = True):
+    def __init__(
+        self,
+        *,
+        components: list[str],
+        atoms_slice: list[int] | SliceLike | None = None,
+        include_ideal_gas: bool = True,
+    ):
         """
         Initialise the observable from a symbolic str component.
 
@@ -258,6 +178,8 @@ def __init__(self, components: list[str], *, include_ideal_gas: bool = True):
         ----------
         components : list[str]
             Symbols for tensor components, xx, yy, etc.
+        atoms_slice : list[int] | SliceLike | None = None
+            List or slice of atoms to observe velocities from.
         include_ideal_gas : bool
             Calculate with the ideal gas contribution.
         """
@@ -277,6 +199,11 @@ def __init__(self, components: list[str], *, include_ideal_gas: bool = True):
         )
         self._set_components(components)
 
+        if atoms_slice:
+            self.atoms_slice = atoms_slice
+        else:
+            self.atoms_slice = slice(0, None, 1)
+
         Observable.__init__(self, len(components))
         self.include_ideal_gas = include_ideal_gas
 
@@ -298,14 +225,18 @@ def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
         list[float]
             The stress components in GPa units.
         """
+        sliced_atoms = atoms[self.atoms_slice]
+        sliced_atoms.calc = atoms.calc
         return (
-            atoms.get_stress(include_ideal_gas=self.include_ideal_gas, voigt=True)
+            sliced_atoms.get_stress(
+                include_ideal_gas=self.include_ideal_gas, voigt=True
+            )
             / units.GPa
         )[self._indices]
 
 
-StressDiagonal = Stress(["xx", "yy", "zz"])
-ShearStress = Stress(["xy", "yz", "zx"])
+StressDiagonal = Stress(components=["xx", "yy", "zz"])
+ShearStress = Stress(components=["xy", "yz", "zx"])
 
 
 # pylint: disable=too-few-public-methods
@@ -317,14 +248,15 @@ class Velocity(Observable, ComponentMixin):
     ----------
     components : list[str]
         Symbols for velocity components, x, y, z.
-    atoms : Optional[Union[list[int], SliceLike]]
+    atoms_slice : list[int] | SliceLike | None = None
         List or slice of atoms to observe velocities from.
     """
 
     def __init__(
         self,
+        *,
         components: list[str],
-        atoms: list[int] | SliceLike | None = None,
+        atoms_slice: list[int] | SliceLike | None = None,
     ):
         """
         Initialise the observable from a symbolic str component and atom index.
@@ -333,17 +265,18 @@ def __init__(
         ----------
         components : list[str]
             Symbols for tensor components, x, y, and z.
-        atoms : Union[list[int], SliceLike]
+        atoms_slice : Union[list[int], SliceLike]
             List or slice of atoms to observe velocities from.
         """
         ComponentMixin.__init__(self, components={"x": 0, "y": 1, "z": 2})
         self._set_components(components)
 
         Observable.__init__(self, len(components))
-        if atoms:
-            self.atoms = atoms
+
+        if atoms_slice:
+            self.atoms_slice = atoms_slice
         else:
-            atoms = slice(None, None, None)
+            self.atoms_slice = slice(0, None, 1)
 
     def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
         """
@@ -363,4 +296,22 @@ def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
         list[float]
             The velocity values.
         """
-        return atoms.get_velocities()[self.atoms, :][:, self._indices].flatten()
+        return atoms.get_velocities()[self.atoms_slice, :][:, self._indices].flatten()
+
+    def value_count(self, n_atoms: int | None = None) -> int:
+        """
+        Count of values returned by __call__.
+
+        Parameters
+        ----------
+        n_atoms : int | None
+            Atom count to expand atoms_slice.
+
+        Returns
+        -------
+        int
+            The number of values returned by __call__.
+        """
+        if isinstance(self.atoms_slice, list):
+            return len(self.atoms_slice) * self.dimension
+        return slicelike_len_for(self.atoms_slice, self.n_atoms) * self.dimension
diff --git a/tests/test_correlator.py b/tests/test_correlator.py
index a6c5f03e..a8c8b40c 100644
--- a/tests/test_correlator.py
+++ b/tests/test_correlator.py
@@ -96,8 +96,8 @@ def test_vaf(tmp_path):
         file_prefix=file_prefix,
         correlation_kwargs=[
             {
-                "a": Velocity(["x", "y", "z"], na),
-                "b": Velocity(["x", "y", "z"], na),
+                "a": Velocity(components=["x", "y", "z"], atoms_slice=na),
+                "b": Velocity(components=["x", "y", "z"], atoms_slice=na),
                 "name": "vaf_Na",
                 "blocks": 1,
                 "points": 11,
@@ -105,8 +105,8 @@ def test_vaf(tmp_path):
                 "update_frequency": 1,
             },
             {
-                "a": Velocity(["x", "y", "z"], cl),
-                "b": Velocity(["x", "y", "z"], cl),
+                "a": Velocity(components=["x", "y", "z"], atoms_slice=cl),
+                "b": Velocity(components=["x", "y", "z"], atoms_slice=cl),
                 "name": "vaf_Cl",
                 "blocks": 1,
                 "points": 11,
@@ -130,8 +130,8 @@ def test_vaf(tmp_path):
         vaf = safe_load(cor)
     vaf_na = np.array(vaf["vaf_Na"]["value"])
     vaf_cl = np.array(vaf["vaf_Cl"]["value"])
-    assert vaf_na == approx(vaf_post[0], rel=1e-5)
-    assert vaf_cl == approx(vaf_post[1], rel=1e-5)
+    assert vaf_na * 3 == approx(vaf_post[0], rel=1e-5)
+    assert vaf_cl * 3 == approx(vaf_post[1], rel=1e-5)
 
 
 def test_md_correlations(tmp_path):
@@ -156,8 +156,8 @@ def test_md_correlations(tmp_path):
         file_prefix=file_prefix,
         correlation_kwargs=[
             {
-                "a": Stress([("xy")]),
-                "b": Stress([("xy")]),
+                "a": Stress(components=[("xy")]),
+                "b": Stress(components=[("xy")]),
                 "name": "stress_xy_auto_cor",
                 "blocks": 1,
                 "points": 11,

From e5ce55d277d8228bf63b47b9db276b980b755ee0 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Fri, 1 Nov 2024 13:33:09 +0000
Subject: [PATCH 08/37] Remove args and kwargs

---
 janus_core/helpers/janus_types.py    |  8 +++---
 janus_core/processing/correlator.py  | 37 +++++++++++-----------------
 janus_core/processing/observables.py | 18 +++-----------
 3 files changed, 21 insertions(+), 42 deletions(-)

diff --git a/janus_core/helpers/janus_types.py b/janus_core/helpers/janus_types.py
index d000ebbb..ae139141 100644
--- a/janus_core/helpers/janus_types.py
+++ b/janus_core/helpers/janus_types.py
@@ -159,10 +159,10 @@ class EoSResults(TypedDict, total=False):
 class CorrelationKwargs(TypedDict, total=True):
     """Arguments for on-the-fly correlations <ab>."""
 
-    #: observable a in <ab>, with optional args and kwargs
-    a: Union[Observable, tuple[Observable, tuple, dict]]
-    #: observable b in <ab>, with optional args and kwargs
-    b: Union[Observable, tuple[Observable, tuple, dict]]
+    #: observable a in <ab>
+    a: Observable
+    #: observable b in <ab>
+    b: Observable
     #: name used for correlation in output
     name: str
     #: blocks used in multi-tau algorithm
diff --git a/janus_core/processing/correlator.py b/janus_core/processing/correlator.py
index d9e9eb7d..3bdbbde1 100644
--- a/janus_core/processing/correlator.py
+++ b/janus_core/processing/correlator.py
@@ -181,10 +181,10 @@ class Correlation:
     ----------
     n_atoms : int
         Number of possible atoms to track.
-    a : Union[Observable, tuple[Observable, tuple, dict]]
-        Getter for a and kwargs.
-    b : Union[Observable, tuple[Observable, tuple, dict]]
-        Getter for b and kwargs.
+    a : Observable
+        Observable for a.
+    b : Observable
+        Observable for b.
     name : str
         Name of correlation.
     blocks : int
@@ -201,8 +201,8 @@ def __init__(
         self,
         *,
         n_atoms: int,
-        a: Observable | tuple[Observable, tuple, dict],
-        b: Observable | tuple[Observable, tuple, dict],
+        a: Observable,
+        b: Observable,
         name: str,
         blocks: int,
         points: int,
@@ -216,10 +216,10 @@ def __init__(
         ----------
         n_atoms : int
             Number of possible atoms to track.
-        a : Union[Observable, tuple[Observable, tuple, dict]]
-            Getter for a and kwargs.
-        b : Union[Observable, tuple[Observable, tuple, dict]]
-            Getter for b and kwargs.
+        a : Observable
+            Observable for a.
+        b : Observable
+            Observable for b.
         name : str
             Name of correlation.
         blocks : int
@@ -232,17 +232,8 @@ def __init__(
             Frequency to update the correlation, md steps.
         """
         self.name = name
-        if isinstance(a, tuple):
-            self._get_a, self._a_args, self._a_kwargs = a
-        else:
-            self._get_a = a
-            self._a_args, self._a_kwargs = (), {}
-
-        if isinstance(b, tuple):
-            self._get_b, self._b_args, self._b_kwargs = b
-        else:
-            self._get_b = b
-            self._b_args, self._b_kwargs = (), {}
+        self._get_a = a
+        self._get_b = b
 
         a_values = self._get_a.value_count(n_atoms)
         b_values = self._get_b.value_count(n_atoms)
@@ -281,8 +272,8 @@ def update(self, atoms: Atoms) -> None:
         """
         for i, values in enumerate(
             zip(
-                self._get_a(atoms, *self._a_args, **self._a_kwargs),
-                self._get_b(atoms, *self._b_args, **self._b_kwargs),
+                self._get_a(atoms),
+                self._get_b(atoms),
             )
         ):
             self._correlators[i].update(*values)
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index d17b8b2b..223fdcf3 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -33,7 +33,7 @@ def __init__(self, dimension: int = 1):
         """
         self._dimension = dimension
 
-    def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
+    def __call__(self, atoms: Atoms) -> list[float]:
         """
         Signature for returning observed value from atoms.
 
@@ -41,10 +41,6 @@ def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
         ----------
         atoms : Atoms
             Atoms object to extract values from.
-        *args : tuple
-            Additional positional arguments passed to getter.
-        **kwargs : dict
-            Additional kwargs passed getter.
 
         Returns
         -------
@@ -207,7 +203,7 @@ def __init__(
         Observable.__init__(self, len(components))
         self.include_ideal_gas = include_ideal_gas
 
-    def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
+    def __call__(self, atoms: Atoms) -> list[float]:
         """
         Get the stress components.
 
@@ -215,10 +211,6 @@ def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
         ----------
         atoms : Atoms
             Atoms object to extract values from.
-        *args : tuple
-            Additional positional arguments passed to getter.
-        **kwargs : dict
-            Additional kwargs passed getter.
 
         Returns
         -------
@@ -278,7 +270,7 @@ def __init__(
         else:
             self.atoms_slice = slice(0, None, 1)
 
-    def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
+    def __call__(self, atoms: Atoms) -> list[float]:
         """
         Get the velocity components for correlated atoms.
 
@@ -286,10 +278,6 @@ def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
         ----------
         atoms : Atoms
             Atoms object to extract values from.
-        *args : tuple
-            Additional positional arguments passed to getter.
-        **kwargs : dict
-            Additional kwargs passed getter.
 
         Returns
         -------

From a1a889bcd75d472186c003591ff21f38ad1e3bad Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Fri, 1 Nov 2024 13:46:19 +0000
Subject: [PATCH 09/37] typo

---
 janus_core/processing/observables.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 223fdcf3..bc245a5a 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -105,7 +105,7 @@ def allowed_components(self) -> dict[str, int]:
 
         Returns
         -------
-        Dict[str, int]
+        dict[str, int]
             The allowed components and associated indices.
         """
         return self._components

From 0726349e1b8a930e1ad1d98a61f719a7b509e3fb Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Fri, 1 Nov 2024 14:32:20 +0000
Subject: [PATCH 10/37] Manually set __module__

---
 janus_core/processing/observables.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index bc245a5a..9f5ca830 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -147,6 +147,7 @@ def _set_components(self, components: list[str]):
 
 # pylint: disable=too-few-public-methods
 class Stress(Observable, ComponentMixin):
+    __module__ = "observables"
     """
     Observable for stress components.
 
@@ -233,6 +234,7 @@ def __call__(self, atoms: Atoms) -> list[float]:
 
 # pylint: disable=too-few-public-methods
 class Velocity(Observable, ComponentMixin):
+    __module__ = "observables"
     """
     Observable for per atom velocity components.
 

From c30d9238296581e27900e1f239bfcbb443d62b78 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Fri, 1 Nov 2024 14:40:30 +0000
Subject: [PATCH 11/37] move to fix pre-commit

---
 janus_core/processing/observables.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 9f5ca830..6b82239d 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -147,7 +147,6 @@ def _set_components(self, components: list[str]):
 
 # pylint: disable=too-few-public-methods
 class Stress(Observable, ComponentMixin):
-    __module__ = "observables"
     """
     Observable for stress components.
 
@@ -161,6 +160,8 @@ class Stress(Observable, ComponentMixin):
         Calculate with the ideal gas contribution.
     """
 
+    __module__ = "observables"
+
     def __init__(
         self,
         *,
@@ -234,7 +235,6 @@ def __call__(self, atoms: Atoms) -> list[float]:
 
 # pylint: disable=too-few-public-methods
 class Velocity(Observable, ComponentMixin):
-    __module__ = "observables"
     """
     Observable for per atom velocity components.
 
@@ -246,6 +246,8 @@ class Velocity(Observable, ComponentMixin):
         List or slice of atoms to observe velocities from.
     """
 
+    __module__ = "observables"
+
     def __init__(
         self,
         *,

From ffeca39845160972b4ba4ebfa5f38e50526087d6 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <33522054+harveydevereux@users.noreply.github.com>
Date: Fri, 1 Nov 2024 14:46:54 +0000
Subject: [PATCH 12/37] Apply suggestions from code review

Apply suggestions

Co-authored-by: Jacob Wilkins <46597752+oerc0122@users.noreply.github.com>
---
 janus_core/processing/correlator.py | 19 ++++++-------------
 1 file changed, 6 insertions(+), 13 deletions(-)

diff --git a/janus_core/processing/correlator.py b/janus_core/processing/correlator.py
index 3bdbbde1..f93ee521 100644
--- a/janus_core/processing/correlator.py
+++ b/janus_core/processing/correlator.py
@@ -242,11 +242,8 @@ def __init__(
             raise ValueError("Observables have inconsistent sizes")
         self._values = a_values
 
-        self._correlators = []
-        for _ in range(self._values):
-            self._correlators.append(
-                Correlator(blocks=blocks, points=points, averaging=averaging)
-            )
+        self._correlators = [Correlator(blocks=blocks, points=points, averaging=averaging)
+                             for _ in range(self._values)]
         self._update_frequency = update_frequency
 
     @property
@@ -270,13 +267,9 @@ def update(self, atoms: Atoms) -> None:
         atoms : Atoms
             Atoms object to observe values from.
         """
-        for i, values in enumerate(
-            zip(
-                self._get_a(atoms),
-                self._get_b(atoms),
-            )
-        ):
-            self._correlators[i].update(*values)
+        atom_pairs = zip(self._get_a(atoms), self._get_b(atoms))
+        for corr, values in zip(self._correlators, atom_pairs):
+            corr.update(*values)
 
     def get(self) -> tuple[Iterable[float], Iterable[float]]:
         """
@@ -291,7 +284,7 @@ def get(self) -> tuple[Iterable[float], Iterable[float]]:
         """
         if self._correlators:
             _, lags = self._correlators[0].get()
-            avg_value = sum([cor.get()[0] for cor in self._correlators]) / self._values
+            avg_value = np.mean(cor.get()[0] for cor in self._correlators)
             return avg_value, lags
         return [], []
 

From 518736d064c42f45e3636728b39fc62517f7de19 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Fri, 1 Nov 2024 15:32:49 +0000
Subject: [PATCH 13/37] abstract method, fix averaging

---
 janus_core/processing/correlator.py  | 9 +++++----
 janus_core/processing/observables.py | 4 +++-
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/janus_core/processing/correlator.py b/janus_core/processing/correlator.py
index f93ee521..8e3ff60d 100644
--- a/janus_core/processing/correlator.py
+++ b/janus_core/processing/correlator.py
@@ -242,8 +242,10 @@ def __init__(
             raise ValueError("Observables have inconsistent sizes")
         self._values = a_values
 
-        self._correlators = [Correlator(blocks=blocks, points=points, averaging=averaging)
-                             for _ in range(self._values)]
+        self._correlators = [
+            Correlator(blocks=blocks, points=points, averaging=averaging)
+            for _ in range(self._values)
+        ]
         self._update_frequency = update_frequency
 
     @property
@@ -284,8 +286,7 @@ def get(self) -> tuple[Iterable[float], Iterable[float]]:
         """
         if self._correlators:
             _, lags = self._correlators[0].get()
-            avg_value = np.mean(cor.get()[0] for cor in self._correlators)
-            return avg_value, lags
+            return np.mean([cor.get()[0] for cor in self._correlators], axis=0), lags
         return [], []
 
     def __str__(self) -> str:
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 6b82239d..5f1404d6 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
 from ase import Atoms, units
@@ -12,7 +13,7 @@
 
 
 # pylint: disable=too-few-public-methods
-class Observable:
+class Observable(ABC):
     """
     Observable data that may be correlated.
 
@@ -33,6 +34,7 @@ def __init__(self, dimension: int = 1):
         """
         self._dimension = dimension
 
+    @abstractmethod
     def __call__(self, atoms: Atoms) -> list[float]:
         """
         Signature for returning observed value from atoms.

From 3c486b5aa3ff1381eb5f1bdec0e783c13e1cd97d Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Mon, 4 Nov 2024 09:44:03 +0000
Subject: [PATCH 14/37] fix import

---
 janus_core/helpers/janus_types.py   | 2 +-
 janus_core/processing/correlator.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/janus_core/helpers/janus_types.py b/janus_core/helpers/janus_types.py
index ae139141..2dacee0e 100644
--- a/janus_core/helpers/janus_types.py
+++ b/janus_core/helpers/janus_types.py
@@ -13,7 +13,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from janus_core.helpers.observables import Observable
+from janus_core.processing.observables import Observable
 
 # General
 
diff --git a/janus_core/processing/correlator.py b/janus_core/processing/correlator.py
index 8e3ff60d..d612cdeb 100644
--- a/janus_core/processing/correlator.py
+++ b/janus_core/processing/correlator.py
@@ -7,7 +7,7 @@
 from ase import Atoms
 import numpy as np
 
-from janus_core.helpers.observables import Observable
+from janus_core.processing.observables import Observable
 
 
 class Correlator:

From 143d082ee31d143dc292cf583081f76de0fa1873 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Mon, 4 Nov 2024 09:52:04 +0000
Subject: [PATCH 15/37] fix import

---
 tests/test_utils.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/test_utils.py b/tests/test_utils.py
index 03cd2915..b8f9bbcb 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -12,10 +12,8 @@
 from janus_core.helpers.janus_types import SliceLike, StartStopStep
 from janus_core.helpers.mlip_calculators import choose_calculator
 from janus_core.helpers.struct_io import output_structs
-from janus_core.helpers.utils import none_to_dict
 from janus_core.helpers.utils import (
     none_to_dict,
-    output_structs,
     slicelike_len_for,
     slicelike_to_startstopstep,
 )

From 34f57eb915550b6bd445813fdf87fa26b44cdd46 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Mon, 4 Nov 2024 10:01:25 +0000
Subject: [PATCH 16/37] ignore unfound refs

---
 docs/source/conf.py                  |  4 ++++
 janus_core/helpers/utils.py          | 14 +++++++++-----
 janus_core/processing/observables.py |  4 ----
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/docs/source/conf.py b/docs/source/conf.py
index 1ca0dcbb..5cc1ca8b 100755
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -203,9 +203,13 @@
     ("py:class", "Architectures"),
     ("py:class", "Devices"),
     ("py:class", "MaybeSequence"),
+    ("py:class", "SliceLike"),
     ("py:class", "PathLike"),
     ("py:class", "Atoms"),
     ("py:class", "Calculator"),
     ("py:class", "Context"),
     ("py:class", "Path"),
+    ("py:obj", "dimension"),
+    ("py:obj", "allowed_components"),
+    ("py:obj", "janus_core.processing.observables.Stress.value_count")
 ]
diff --git a/janus_core/helpers/utils.py b/janus_core/helpers/utils.py
index f4d28101..ea096896 100644
--- a/janus_core/helpers/utils.py
+++ b/janus_core/helpers/utils.py
@@ -18,7 +18,13 @@
 )
 from rich.style import Style
 
-from janus_core.helpers.janus_types import MaybeSequence, PathLike, SliceLike, StartStopStep
+from janus_core.helpers.janus_types import (
+    MaybeSequence,
+    PathLike,
+    SliceLike,
+    StartStopStep,
+)
+
 
 class FileNameMixin(ABC):  # noqa: B024 (abstract-base-class-without-abstract-method)
     """
@@ -409,6 +415,7 @@ def track_progress(sequence: Sequence | Iterable, description: str) -> Iterable:
     with progress:
         yield from progress.track(sequence, description=description)
 
+
 def slicelike_to_startstopstep(index: SliceLike) -> StartStopStep:
     """
     Standarize `SliceLike`s into tuple of `start`, `stop`, `step`.
@@ -453,7 +460,4 @@ def slicelike_len_for(slc: SliceLike, sliceable_length: int) -> int:
     start, stop, step = slicelike_to_startstopstep(slc)
     if stop is None:
         stop = sliceable_length
-    # start = start if start is None else 0
-    # stop = stop if stop is None else sliceable_length
-    # step = step if step is None else 1
-    return len(range(start, stop, step))
\ No newline at end of file
+    return len(range(start, stop, step))
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 5f1404d6..07a0626f 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -162,8 +162,6 @@ class Stress(Observable, ComponentMixin):
         Calculate with the ideal gas contribution.
     """
 
-    __module__ = "observables"
-
     def __init__(
         self,
         *,
@@ -248,8 +246,6 @@ class Velocity(Observable, ComponentMixin):
         List or slice of atoms to observe velocities from.
     """
 
-    __module__ = "observables"
-
     def __init__(
         self,
         *,

From b41034fd05302129cc3c80dd8b15c08965563fa5 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Mon, 4 Nov 2024 10:07:12 +0000
Subject: [PATCH 17/37] move import

---
 tests/test_correlator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_correlator.py b/tests/test_correlator.py
index a8c8b40c..6521d7b0 100644
--- a/tests/test_correlator.py
+++ b/tests/test_correlator.py
@@ -14,9 +14,9 @@
 
 from janus_core.calculations.md import NVE
 from janus_core.calculations.single_point import SinglePoint
+from janus_core.processing import post_process
 from janus_core.processing.correlator import Correlator
 from janus_core.processing.observables import Stress, Velocity
-from janus_core.processing import post_process
 
 DATA_PATH = Path(__file__).parent / "data"
 MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model"

From 2e82eba95cd6b0e4fd55da89c6c0c58c73c8d4f4 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Mon, 4 Nov 2024 15:26:51 +0000
Subject: [PATCH 18/37] rebase for vaf lags

---
 tests/test_correlator.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_correlator.py b/tests/test_correlator.py
index 6521d7b0..608602a9 100644
--- a/tests/test_correlator.py
+++ b/tests/test_correlator.py
@@ -130,8 +130,8 @@ def test_vaf(tmp_path):
         vaf = safe_load(cor)
     vaf_na = np.array(vaf["vaf_Na"]["value"])
     vaf_cl = np.array(vaf["vaf_Cl"]["value"])
-    assert vaf_na * 3 == approx(vaf_post[0], rel=1e-5)
-    assert vaf_cl * 3 == approx(vaf_post[1], rel=1e-5)
+    assert vaf_na * 3 == approx(vaf_post[1][0], rel=1e-5)
+    assert vaf_cl * 3 == approx(vaf_post[1][1], rel=1e-5)
 
 
 def test_md_correlations(tmp_path):

From 773598cb953944dd7667d466d4851427bc066ac2 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <33522054+harveydevereux@users.noreply.github.com>
Date: Mon, 4 Nov 2024 16:02:26 +0000
Subject: [PATCH 19/37] Apply suggestions from code review

apply review comments

Co-authored-by: Jacob Wilkins <46597752+oerc0122@users.noreply.github.com>
---
 janus_core/processing/observables.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 07a0626f..f2a12f8b 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -138,9 +138,8 @@ def _set_components(self, components: list[str]):
         ValueError
             If any component is invalid.
         """
-        for component in components:
+        for component in self.allowed_components.keys() - components.keys():
             if component not in self.allowed_components:
-                component_names = list(self._components.keys())
                 raise ValueError(
                     f"'{component}' invalid, must be '{', '.join(component_names)}'"
                 )
@@ -267,10 +266,7 @@ def __init__(
 
         Observable.__init__(self, len(components))
 
-        if atoms_slice:
-            self.atoms_slice = atoms_slice
-        else:
-            self.atoms_slice = slice(0, None, 1)
+        self.atoms_slice = atoms_slice if atoms_slice else slice(0, None, 1)
 
     def __call__(self, atoms: Atoms) -> list[float]:
         """

From e2a46a6fa4d82ffd6509be31fb86a3dc50b3397b Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Mon, 4 Nov 2024 16:41:50 +0000
Subject: [PATCH 20/37] Rename builtins, multi-line error msg

---
 janus_core/processing/observables.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index f2a12f8b..1aee858d 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -138,11 +138,12 @@ def _set_components(self, components: list[str]):
         ValueError
             If any component is invalid.
         """
-        for component in self.allowed_components.keys() - components.keys():
-            if component not in self.allowed_components:
-                raise ValueError(
-                    f"'{component}' invalid, must be '{', '.join(component_names)}'"
-                )
+        if any(components - self.allowed_components.keys()):
+            raise ValueError(
+                f"'{components-self.allowed_components.keys()}'"
+                " invalid, must be '{', '.join(self._components)}'"
+            )
+
         self.components = components
 
 
@@ -228,8 +229,8 @@ def __call__(self, atoms: Atoms) -> list[float]:
         )[self._indices]
 
 
-StressDiagonal = Stress(components=["xx", "yy", "zz"])
-ShearStress = Stress(components=["xy", "yz", "zx"])
+StressHydrostatic = Stress(components=["xx", "yy", "zz"])
+StressShear = Stress(components=["xy", "yz", "zx"])
 
 
 # pylint: disable=too-few-public-methods

From 514d81b9d75e063408ccd132674b0793a40e0547 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Mon, 4 Nov 2024 17:04:47 +0000
Subject: [PATCH 21/37] remove uneeded property

---
 janus_core/processing/observables.py | 18 +++---------------
 1 file changed, 3 insertions(+), 15 deletions(-)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 1aee858d..fe05b3a0 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -98,19 +98,7 @@ def __init__(self, components: dict[str, int]):
         components : dict[str, int]
             Symbolic components mapped to indices.
         """
-        self._components = components
-
-    @property
-    def allowed_components(self) -> dict[str, int]:
-        """
-        Allowed symbolic components with associated indices.
-
-        Returns
-        -------
-        dict[str, int]
-            The allowed components and associated indices.
-        """
-        return self._components
+        self._allowed_components = components
 
     @property
     def _indices(self) -> list[int]:
@@ -122,7 +110,7 @@ def _indices(self) -> list[int]:
         list[int]
             The indices for each self._components.
         """
-        return [self._components[c] for c in self.components]
+        return [self._allowed_components[c] for c in self.components]
 
     def _set_components(self, components: list[str]):
         """
@@ -138,7 +126,7 @@ def _set_components(self, components: list[str]):
         ValueError
             If any component is invalid.
         """
-        if any(components - self.allowed_components.keys()):
+        if any(components - self._allowed_components.keys()):
             raise ValueError(
                 f"'{components-self.allowed_components.keys()}'"
                 " invalid, must be '{', '.join(self._components)}'"

From 18e5f3579ce2372056f331b009bdbf9d650af1e0 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Tue, 5 Nov 2024 10:39:35 +0000
Subject: [PATCH 22/37] restore spacing

---
 janus_core/helpers/janus_types.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/janus_core/helpers/janus_types.py b/janus_core/helpers/janus_types.py
index 2dacee0e..7edad5ad 100644
--- a/janus_core/helpers/janus_types.py
+++ b/janus_core/helpers/janus_types.py
@@ -23,6 +23,7 @@
 PathLike = Union[str, Path]
 StartStopStep = tuple[Optional[int], Optional[int], int]
 SliceLike = Union[slice, range, int, StartStopStep]
+
 # ASE Arg types
 
 

From aba3aeaeb50cdc177fbeb9dab2fabb4a7c1d955c Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Tue, 5 Nov 2024 11:05:28 +0000
Subject: [PATCH 23/37] Update developer guide

---
 docs/source/developer_guide/tutorial.rst | 28 ++++++++++++++++++------
 1 file changed, 21 insertions(+), 7 deletions(-)

diff --git a/docs/source/developer_guide/tutorial.rst b/docs/source/developer_guide/tutorial.rst
index e2154ccf..269b485b 100644
--- a/docs/source/developer_guide/tutorial.rst
+++ b/docs/source/developer_guide/tutorial.rst
@@ -187,20 +187,34 @@ Alternatively, using ``tox``::
 Adding a new Observable
 =======================
 
-Additional built-in observable quantities may be added for use by the ``janus_core.processing.correlator.Correlation`` class. These should conform to the ``__call__`` signature of ``janus_core.helpers.janus_types.Observable``. For a user this can be accomplished by writing a function, or class also implementing a commensurate ``__call__``.
+Additional built-in observable quantities may be added for use by the ``janus_core.processing.correlator.Correlation`` class. These should extend ``janus_core.processing.observables.Observable``. The abstract method ``__call__`` should be implemented to obtain the values of the observed quantity from an ``Atoms`` object:
 
-Built-in observables are collected within the ``janus_core.processing.observables`` module. For example the ``janus_core.processing.observables.Stress`` observable allows a user to quickly setup a given correlation of stress tensor components (with and without the ideal gas contribution). An observable for the ``xy`` component is obtained without the ideal gas contribution as:
+.. code-block:: python
+
+    @abstractmethod
+    def __call__(self, atoms: Atoms) -> list[float]
+
+These values are returned as a ``list[float]`` representing the dimensions of the observed value which may be returned for a slice of atoms individually. ``value_count`` should be overloaded to return the expected length returned by ``__call__``. This is so enough correlators may be spawned to track the correlated the values. Note the final correlation will be the average of the correlations across these dimensions and atom counts.
+
+For example the ``janus_core.processing.observables.Stress`` built-in returns the required stress components as calculated by ``Atoms.get_stress`` called upon the tracked atoms. Therefore ``__call__`` returns one ``float`` per component correlated. For example an observable representing the ``xy`` and ``zy`` components of stress computed across all odd-index atoms in some ``Atoms`` object can be constructed as follows:
 
 .. code-block:: python
 
-    Stress("xy", False)
+    s = Stress(components=["xy", "zy"], atoms_slice=(0, None, 2))
 
-A new built-in observables can be implemented by a class with the method:
+In this case ``s(atoms)`` will return 2 values.
+
+The ``janus_core.processing.observables.Velocity`` built-in's ``__call__`` not only returns atom velocity for the requested dimensions, but also returns them for every tracked atom. ``value_count`` is overloaded to reflect this. Therefore given the observable
 
 .. code-block:: python
 
-   def __call__(self, atoms: Atoms, *args, **kwargs) -> float
+    v = Velocity(components=["x", "y", "z"], atoms_slice=(0, None, 2))
+
+The value of ``v(atoms)`` will have the length ``3 * len(atoms)//2``. The 3 dimensions for each odd-indexed atom.
 
-The ``__call__`` should contain all the logic for obtaining some ``float`` value from an ``Atoms`` object, alongside optional positional arguments and kwargs. The args and kwargs are set by a user when specifying correlations for a ``janus_core.calculations.md.MolecularDynamics`` run. See also ``janus_core.helpers.janus_types.CorrelationKwargs``. These are set at the instantiation of the ``janus_core.calculations.md.MolecularDynamics`` object and are not modified. These could be used e.g. to specify an observable calculated only from one atom's data.
+New built-in observables are collected within the ``janus_core.processing.observables`` module. Special cases may also be defined for ease of use:
+
+.. code-block:: python
 
-``janus_core.processing.observables.Stress`` includes a constructor to take a symbolic component, e.g. ``"xx"`` or ``"yz"``, and determine the index required from ``ase.Atoms.get_stress`` on instantiation for ease of use.
+    StressHydrostatic = Stress(components=["xx", "yy", "zz"])
+    StressShear = Stress(components=["xy", "yz", "zx"])

From 8c0e9c504634f4cf5bfae1ca9e1ba91195fcdc89 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Tue, 5 Nov 2024 11:30:10 +0000
Subject: [PATCH 24/37] fix error msg

---
 janus_core/processing/observables.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index fe05b3a0..81ecf343 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -128,8 +128,8 @@ def _set_components(self, components: list[str]):
         """
         if any(components - self._allowed_components.keys()):
             raise ValueError(
-                f"'{components-self.allowed_components.keys()}'"
-                " invalid, must be '{', '.join(self._components)}'"
+                f"'{components-self._allowed_components.keys()}'"
+                f" invalid, must be '{', '.join(self._allowed_components)}'"
             )
 
         self.components = components

From 0f5b64862134e55458bd1fd1f20efdc44d289971 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Tue, 5 Nov 2024 12:23:12 +0000
Subject: [PATCH 25/37] fix len_for

---
 janus_core/helpers/utils.py          | 18 +++++++++++-------
 janus_core/processing/observables.py |  6 ++----
 tests/test_utils.py                  | 12 ++++++++----
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/janus_core/helpers/utils.py b/janus_core/helpers/utils.py
index ea096896..b0a93e4c 100644
--- a/janus_core/helpers/utils.py
+++ b/janus_core/helpers/utils.py
@@ -441,23 +441,27 @@ def slicelike_to_startstopstep(index: SliceLike) -> StartStopStep:
     return index
 
 
-def slicelike_len_for(slc: SliceLike, sliceable_length: int) -> int:
+def selector_len(slc: Union[SliceLike, list], selectable_length: int) -> int:
     """
-    Calculate the length of a SliceLike applied to a sliceable of a given length.
+    Calculate the length of a selector applied to an indexable of a given length.
 
     Parameters
     ----------
-    slc : SliceLike
-        The applied SliceLike.
-    sliceable_length : int
-        The length of the sliceable object.
+    slc : Union[SliceLike, list]
+        The applied SliceLike or list for selection.
+    selectable_length : int
+        The length of the selectable object.
 
     Returns
     -------
     int
         Length of the result of applying slc.
     """
+    if isinstance(slc, int):
+        return 1
+    if isinstance(slc, list):
+        return len(slc)
     start, stop, step = slicelike_to_startstopstep(slc)
     if stop is None:
-        stop = sliceable_length
+        stop = selectable_length
     return len(range(start, stop, step))
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 81ecf343..bc6a5fc0 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -9,7 +9,7 @@
 
 if TYPE_CHECKING:
     from janus_core.helpers.janus_types import SliceLike
-    from janus_core.helpers.utils import slicelike_len_for
+    from janus_core.helpers.utils import selector_len
 
 
 # pylint: disable=too-few-public-methods
@@ -287,6 +287,4 @@ def value_count(self, n_atoms: int | None = None) -> int:
         int
             The number of values returned by __call__.
         """
-        if isinstance(self.atoms_slice, list):
-            return len(self.atoms_slice) * self.dimension
-        return slicelike_len_for(self.atoms_slice, self.n_atoms) * self.dimension
+        return selector_len(self.atoms_slice, self.n_atoms) * self.dimension
diff --git a/tests/test_utils.py b/tests/test_utils.py
index b8f9bbcb..916da4a3 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -14,7 +14,7 @@
 from janus_core.helpers.struct_io import output_structs
 from janus_core.helpers.utils import (
     none_to_dict,
-    slicelike_len_for,
+    selector_len,
     slicelike_to_startstopstep,
 )
 
@@ -188,11 +188,15 @@ def test_slicelike_to_startstopstep(slc: SliceLike, expected: StartStopStep):
         ((1, 1), 1),
         ((range(1, 2, 3), 3), 1),
         ((slice(1, 2, 3), 3), 1),
-        ((-1, 1), 2),
+        ((-1, 5), 1),
+        ((-3, 4), 1),
         ((range(10), 10), 10),
         ((slice(0, None, 2), 10), 5),
+        (([-2, -1, 0], 9), 3),
+        (([-1], 10), 1),
+        (([0, -1, 2, 9], 10), 4),
     ],
 )
-def test_slicelike_len_for(slc_len: tuple[SliceLike, int], expected: int):
+def test_selector_len(slc_len: tuple[SliceLike | list, int], expected: int):
     """Test converting SliceLike to StartStopStep."""
-    assert slicelike_len_for(*slc_len) == expected
+    assert selector_len(*slc_len) == expected

From 97444debe6c043c2fa8899bb590a5c36b1764bab Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Tue, 5 Nov 2024 12:32:41 +0000
Subject: [PATCH 26/37] CorrelationKwargs import Observable directly

---
 janus_core/helpers/janus_types.py    | 43 ++++++++++++++--------------
 janus_core/processing/observables.py |  8 ++----
 tests/test_utils.py                  |  3 +-
 3 files changed, 27 insertions(+), 27 deletions(-)

diff --git a/janus_core/helpers/janus_types.py b/janus_core/helpers/janus_types.py
index 7edad5ad..de5c4800 100644
--- a/janus_core/helpers/janus_types.py
+++ b/janus_core/helpers/janus_types.py
@@ -13,8 +13,6 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from janus_core.processing.observables import Observable
-
 # General
 
 T = TypeVar("T")
@@ -98,6 +96,28 @@ class CorrelationKwargs(TypedDict, total=True):
     update_frequency: int
 
 
+class CorrelationKwargs(TypedDict, total=True):
+    """Arguments for on-the-fly correlations <ab>."""
+
+    # imported here to prevent circular imports.
+    from janus_core.processing.observables import Observable
+
+    #: observable a in <ab>, with optional args and kwargs
+    a: Union[Observable, tuple[Observable, tuple, dict]]
+    #: observable b in <ab>, with optional args and kwargs
+    b: Union[Observable, tuple[Observable, tuple, dict]]
+    #: name used for correlation in output
+    name: str
+    #: blocks used in multi-tau algorithm
+    blocks: int
+    #: points per block
+    points: int
+    #: averaging between blocks
+    averaging: int
+    #: frequency to update the correlation (steps)
+    update_frequency: int
+
+
 # eos_names from ase.eos
 EoSNames = Literal[
     "sj",
@@ -155,22 +175,3 @@ class EoSResults(TypedDict, total=False):
     bulk_modulus: float
     v_0: float
     e_0: float
-
-
-class CorrelationKwargs(TypedDict, total=True):
-    """Arguments for on-the-fly correlations <ab>."""
-
-    #: observable a in <ab>
-    a: Observable
-    #: observable b in <ab>
-    b: Observable
-    #: name used for correlation in output
-    name: str
-    #: blocks used in multi-tau algorithm
-    blocks: int
-    #: points per block
-    points: int
-    #: averaging between blocks
-    averaging: int
-    #: frequency to update the correlation (steps)
-    update_frequency: int
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index bc6a5fc0..0f36d202 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -3,13 +3,11 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
 
 from ase import Atoms, units
 
-if TYPE_CHECKING:
-    from janus_core.helpers.janus_types import SliceLike
-    from janus_core.helpers.utils import selector_len
+from janus_core.helpers.janus_types import SliceLike
+from janus_core.helpers.utils import selector_len
 
 
 # pylint: disable=too-few-public-methods
@@ -287,4 +285,4 @@ def value_count(self, n_atoms: int | None = None) -> int:
         int
             The number of values returned by __call__.
         """
-        return selector_len(self.atoms_slice, self.n_atoms) * self.dimension
+        return selector_len(self.atoms_slice, n_atoms) * self.dimension
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 916da4a3..96d810e8 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from typing import Union
 
 from ase import Atoms
 from ase.io import read
@@ -197,6 +198,6 @@ def test_slicelike_to_startstopstep(slc: SliceLike, expected: StartStopStep):
         (([0, -1, 2, 9], 10), 4),
     ],
 )
-def test_selector_len(slc_len: tuple[SliceLike | list, int], expected: int):
+def test_selector_len(slc_len: tuple[Union[SliceLike, list], int], expected: int):
     """Test converting SliceLike to StartStopStep."""
     assert selector_len(*slc_len) == expected

From ef9d3bdbc6d800dfb9ab0790bd2a7457b8fb0b86 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Tue, 5 Nov 2024 15:22:30 +0000
Subject: [PATCH 27/37] Simplify Stress __call__

---
 janus_core/processing/observables.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 0f36d202..82460182 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -207,12 +207,13 @@ def __call__(self, atoms: Atoms) -> list[float]:
         """
         sliced_atoms = atoms[self.atoms_slice]
         sliced_atoms.calc = atoms.calc
-        return (
+        stresses = (
             sliced_atoms.get_stress(
                 include_ideal_gas=self.include_ideal_gas, voigt=True
             )
             / units.GPa
-        )[self._indices]
+        )
+        return stresses[self._indices]
 
 
 StressHydrostatic = Stress(components=["xx", "yy", "zz"])

From 2e7a95d3c4728bfd2f555ffb503b223469323026 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Wed, 6 Nov 2024 11:26:54 +0000
Subject: [PATCH 28/37] Fix typing, use | in janus_types

---
 janus_core/helpers/janus_types.py    | 32 ++++++----------------------
 janus_core/processing/observables.py |  5 ++++-
 2 files changed, 10 insertions(+), 27 deletions(-)

diff --git a/janus_core/helpers/janus_types.py b/janus_core/helpers/janus_types.py
index de5c4800..817c5fc2 100644
--- a/janus_core/helpers/janus_types.py
+++ b/janus_core/helpers/janus_types.py
@@ -6,13 +6,16 @@
 from enum import Enum
 import logging
 from pathlib import Path, PurePath
-from typing import IO, Literal, Optional, TypedDict, TypeVar, Union
+from typing import IO, TYPE_CHECKING, Literal, Optional, TypedDict, TypeVar, Union
 
 from ase import Atoms
 from ase.eos import EquationOfState
 import numpy as np
 from numpy.typing import NDArray
 
+if TYPE_CHECKING:
+    from janus_core.processing.observables import Observable
+
 # General
 
 T = TypeVar("T")
@@ -76,36 +79,13 @@ class PostProcessKwargs(TypedDict, total=False):
     vaf_step: int
     vaf_output_file: PathLike | None
 
-
-class CorrelationKwargs(TypedDict, total=True):
-    """Arguments for on-the-fly correlations <ab>."""
-
-    #: observable a in <ab>, with optional args and kwargs
-    a: Observable | tuple[Observable, tuple, dict]
-    #: observable b in <ab>, with optional args and kwargs
-    b: Observable | tuple[Observable, tuple, dict]
-    #: name used for correlation in output
-    name: str
-    #: blocks used in multi-tau algorithm
-    blocks: int
-    #: points per block
-    points: int
-    #: averaging between blocks
-    averaging: int
-    #: frequency to update the correlation (steps)
-    update_frequency: int
-
-
 class CorrelationKwargs(TypedDict, total=True):
     """Arguments for on-the-fly correlations <ab>."""
 
-    # imported here to prevent circular imports.
-    from janus_core.processing.observables import Observable
-
     #: observable a in <ab>, with optional args and kwargs
-    a: Union[Observable, tuple[Observable, tuple, dict]]
+    a: Observable
     #: observable b in <ab>, with optional args and kwargs
-    b: Union[Observable, tuple[Observable, tuple, dict]]
+    b: Observable
     #: name used for correlation in output
     name: str
     #: blocks used in multi-tau algorithm
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 82460182..9dbf8364 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -3,10 +3,13 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
 
 from ase import Atoms, units
 
-from janus_core.helpers.janus_types import SliceLike
+if TYPE_CHECKING:
+    from janus_core.helpers.janus_types import SliceLike
+
 from janus_core.helpers.utils import selector_len
 
 

From bc3b6ad152676d950ba0e234ee873205d4042b6f Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Wed, 6 Nov 2024 13:31:00 +0000
Subject: [PATCH 29/37] fix atoms_slice, parse in Observable

---
 janus_core/processing/observables.py | 28 +++++++++++++++++-----------
 tests/test_correlator.py             | 20 ++++++++++----------
 2 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 9dbf8364..6969820c 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -10,7 +10,7 @@
 if TYPE_CHECKING:
     from janus_core.helpers.janus_types import SliceLike
 
-from janus_core.helpers.utils import selector_len
+from janus_core.helpers.utils import selector_len, slicelike_to_startstopstep
 
 
 # pylint: disable=too-few-public-methods
@@ -20,20 +20,33 @@ class Observable(ABC):
 
     Parameters
     ----------
+    atoms_slice : list[int] | SliceLike | None = None
+        A slice of atoms to observe.
     dimension : int
         The dimension of the observed data.
     """
 
-    def __init__(self, dimension: int = 1):
+    def __init__(
+        self, atoms_slice: list[int] | SliceLike | None = None, dimension: int = 1
+    ):
         """
         Initialise an observable with a given dimensionality.
 
         Parameters
         ----------
+        atoms_slice : list[int] | SliceLike | None = None
+            A slice of atoms to observe.
         dimension : int
             The dimension of the observed data.
         """
         self._dimension = dimension
+        if atoms_slice:
+            if isinstance(atoms_slice, list):
+                self.atoms_slice = atoms_slice
+            else:
+                self.atoms_slice = slice(*slicelike_to_startstopstep(atoms_slice))
+        else:
+            self.atoms_slice = slice(0, None, 1)
 
     @abstractmethod
     def __call__(self, atoms: Atoms) -> list[float]:
@@ -186,12 +199,7 @@ def __init__(
         )
         self._set_components(components)
 
-        if atoms_slice:
-            self.atoms_slice = atoms_slice
-        else:
-            self.atoms_slice = slice(0, None, 1)
-
-        Observable.__init__(self, len(components))
+        Observable.__init__(self, atoms_slice, len(components))
         self.include_ideal_gas = include_ideal_gas
 
     def __call__(self, atoms: Atoms) -> list[float]:
@@ -255,9 +263,7 @@ def __init__(
         ComponentMixin.__init__(self, components={"x": 0, "y": 1, "z": 2})
         self._set_components(components)
 
-        Observable.__init__(self, len(components))
-
-        self.atoms_slice = atoms_slice if atoms_slice else slice(0, None, 1)
+        Observable.__init__(self, atoms_slice, len(components))
 
     def __call__(self, atoms: Atoms) -> list[float]:
         """
diff --git a/tests/test_correlator.py b/tests/test_correlator.py
index 608602a9..58d993cb 100644
--- a/tests/test_correlator.py
+++ b/tests/test_correlator.py
@@ -78,13 +78,8 @@ def test_vaf(tmp_path):
         calc_kwargs={"model": MODEL_PATH},
     )
 
-    na = []
-    cl = []
-    for i, atom in enumerate(single_point.struct):
-        if atom.symbol == "Na":
-            na.append(i)
-        else:
-            cl.append(i)
+    na = list(range(0, len(single_point.struct), 2))
+    cl = list(range(1, len(single_point.struct), 2))
 
     nve = NVE(
         struct=single_point.struct,
@@ -96,8 +91,11 @@ def test_vaf(tmp_path):
         file_prefix=file_prefix,
         correlation_kwargs=[
             {
-                "a": Velocity(components=["x", "y", "z"], atoms_slice=na),
-                "b": Velocity(components=["x", "y", "z"], atoms_slice=na),
+                "a": Velocity(components=["x", "y", "z"], atoms_slice=(0, None, 2)),
+                "b": Velocity(
+                    components=["x", "y", "z"],
+                    atoms_slice=range(0, len(single_point.struct), 2),
+                ),
                 "name": "vaf_Na",
                 "blocks": 1,
                 "points": 11,
@@ -106,7 +104,9 @@ def test_vaf(tmp_path):
             },
             {
                 "a": Velocity(components=["x", "y", "z"], atoms_slice=cl),
-                "b": Velocity(components=["x", "y", "z"], atoms_slice=cl),
+                "b": Velocity(
+                    components=["x", "y", "z"], atoms_slice=slice(1, None, 2)
+                ),
                 "name": "vaf_Cl",
                 "blocks": 1,
                 "points": 11,

From 4f1c2c248fe986dadf986898a4bb4896012c2b3e Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Wed, 6 Nov 2024 13:37:08 +0000
Subject: [PATCH 30/37] split slc_len in test_selector_len

---
 tests/test_utils.py | 29 ++++++++++++++---------------
 1 file changed, 14 insertions(+), 15 deletions(-)

diff --git a/tests/test_utils.py b/tests/test_utils.py
index 96d810e8..4fdc1a71 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Union
 
 from ase import Atoms
 from ase.io import read
@@ -183,21 +182,21 @@ def test_slicelike_to_startstopstep(slc: SliceLike, expected: StartStopStep):
 
 
 @pytest.mark.parametrize(
-    "slc_len, expected",
+    "slc, len, expected",
     [
-        (((1, 2, 3), 3), 1),
-        ((1, 1), 1),
-        ((range(1, 2, 3), 3), 1),
-        ((slice(1, 2, 3), 3), 1),
-        ((-1, 5), 1),
-        ((-3, 4), 1),
-        ((range(10), 10), 10),
-        ((slice(0, None, 2), 10), 5),
-        (([-2, -1, 0], 9), 3),
-        (([-1], 10), 1),
-        (([0, -1, 2, 9], 10), 4),
+        ((1, 2, 3), 3, 1),
+        (1, 1, 1),
+        (range(1, 2, 3), 3, 1),
+        (slice(1, 2, 3), 3, 1),
+        (-1, 5, 1),
+        (-3, 4, 1),
+        (range(10), 10, 10),
+        (slice(0, None, 2), 10, 5),
+        ([-2, -1, 0], 9, 3),
+        ([-1], 10, 1),
+        ([0, -1, 2, 9], 10, 4),
     ],
 )
-def test_selector_len(slc_len: tuple[Union[SliceLike, list], int], expected: int):
+def test_selector_len(slc: SliceLike | list[int], len: int, expected: int):
     """Test converting SliceLike to StartStopStep."""
-    assert selector_len(*slc_len) == expected
+    assert selector_len(slc, len) == expected

From 03ce113853947e2bf41da0724453831e7c0a064d Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Wed, 6 Nov 2024 13:44:32 +0000
Subject: [PATCH 31/37] minimal exclusions, fix typing

---
 docs/source/conf.py | 1 -
 tests/test_utils.py | 3 ++-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/source/conf.py b/docs/source/conf.py
index 5cc1ca8b..51483d12 100755
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -210,6 +210,5 @@
     ("py:class", "Context"),
     ("py:class", "Path"),
     ("py:obj", "dimension"),
-    ("py:obj", "allowed_components"),
     ("py:obj", "janus_core.processing.observables.Stress.value_count")
 ]
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 4fdc1a71..4e9baa11 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from typing import Union
 
 from ase import Atoms
 from ase.io import read
@@ -197,6 +198,6 @@ def test_slicelike_to_startstopstep(slc: SliceLike, expected: StartStopStep):
         ([0, -1, 2, 9], 10, 4),
     ],
 )
-def test_selector_len(slc: SliceLike | list[int], len: int, expected: int):
+def test_selector_len(slc: Union[SliceLike, list[int]], len: int, expected: int):
     """Test converting SliceLike to StartStopStep."""
     assert selector_len(slc, len) == expected

From 5c36b3de892b8fb3697055783ada1a8ae63ae1aa Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Wed, 6 Nov 2024 14:44:19 +0000
Subject: [PATCH 32/37] expand on doc

---
 docs/source/developer_guide/tutorial.rst | 42 +++++++++++++++++-------
 1 file changed, 31 insertions(+), 11 deletions(-)

diff --git a/docs/source/developer_guide/tutorial.rst b/docs/source/developer_guide/tutorial.rst
index 269b485b..76f392fc 100644
--- a/docs/source/developer_guide/tutorial.rst
+++ b/docs/source/developer_guide/tutorial.rst
@@ -187,32 +187,52 @@ Alternatively, using ``tox``::
 Adding a new Observable
 =======================
 
-Additional built-in observable quantities may be added for use by the ``janus_core.processing.correlator.Correlation`` class. These should extend ``janus_core.processing.observables.Observable``. The abstract method ``__call__`` should be implemented to obtain the values of the observed quantity from an ``Atoms`` object:
+Additional built-in observable quantities may be added for use by the ``janus_core.processing.correlator.Correlation`` class. These should extend ``janus_core.processing.observables.Observable``. The abstract method ``__call__`` should be implemented to obtain the values of the observed quantity from an ``Atoms`` object, and potentially ``value_count`` may need to be overloaded indicating the number of values returned by ``__call__``. The number of values returned will depend on the dimension of the underlying ``Observable``, and potentially the number of atoms resulting from the ``Observable``'s slice. These values are correlated individually and averaged.
 
-.. code-block:: python
+As an example of building a new ``Observable`` consider the ``janus_core.processing.observables.Stress`` and ``janus_core.processing.observables.Velocity`` built-ins. ``janus_core.processing.observables.Stress`` returns the required stress components as calculated by ``Atoms.get_stress`` called upon the tracked atoms. Therefore ``__call__`` returns one ``float`` per component correlated. The call method applies the atom slice and returns the desired stress components.
 
-    @abstractmethod
-    def __call__(self, atoms: Atoms) -> list[float]
+.. code-block:: python
 
-These values are returned as a ``list[float]`` representing the dimensions of the observed value which may be returned for a slice of atoms individually. ``value_count`` should be overloaded to return the expected length returned by ``__call__``. This is so enough correlators may be spawned to track the correlated the values. Note the final correlation will be the average of the correlations across these dimensions and atom counts.
+    def __call__(self, atoms: Atoms) -> list[float]:
+        sliced_atoms = atoms[self.atoms_slice]
+        sliced_atoms.calc = atoms.calc
+        stresses = (
+            sliced_atoms.get_stress(
+                include_ideal_gas=self.include_ideal_gas, voigt=True
+            )
+            / units.GPa
+        )
+        return stresses[self._indices]
 
-For example the ``janus_core.processing.observables.Stress`` built-in returns the required stress components as calculated by ``Atoms.get_stress`` called upon the tracked atoms. Therefore ``__call__`` returns one ``float`` per component correlated. For example an observable representing the ``xy`` and ``zy`` components of stress computed across all odd-index atoms in some ``Atoms`` object can be constructed as follows:
+For example an observable representing the ``xy`` and ``zy`` components of stress computed across all odd-index atoms in some ``Atoms`` object can be constructed as follows:
 
 .. code-block:: python
 
     s = Stress(components=["xy", "zy"], atoms_slice=(0, None, 2))
 
-In this case ``s(atoms)`` will return 2 values.
+In this case ``s(atoms)`` will return 2 values, no matter how many atoms are sliced from ``atoms_slice``.
 
-The ``janus_core.processing.observables.Velocity`` built-in's ``__call__`` not only returns atom velocity for the requested dimensions, but also returns them for every tracked atom. ``value_count`` is overloaded to reflect this. Therefore given the observable
+However, the ``janus_core.processing.observables.Velocity`` built-in's ``__call__`` not only returns atom velocity for the requested dimensions, but also returns them for every tracked atom, i.e:
 
 .. code-block:: python
 
-    v = Velocity(components=["x", "y", "z"], atoms_slice=(0, None, 2))
+       def __call__(self, atoms: Atoms) -> list[float]:
+           return atoms.get_velocities()[self.atoms_slice, :][:, self._indices].flatten()
+
+so ``value_count`` is overloaded to reflect this.
+
+.. code-block:: python
+
+       def value_count(self, n_atoms: int | None = None) -> int:
+           return selector_len(self.atoms_slice, n_atoms) * self.dimension
+
+Unlike ``janus_core.processing.observables.Stress`` when we create a ``janus_core.processing.observables.Velocity`` observable like so,
 
-The value of ``v(atoms)`` will have the length ``3 * len(atoms)//2``. The 3 dimensions for each odd-indexed atom.
+.. code-block:: python
+
+    v = Velocity(components=["x", "y", "z"], atoms_slice=(0, None, 2))
 
-New built-in observables are collected within the ``janus_core.processing.observables`` module. Special cases may also be defined for ease of use:
+the value of ``v(atoms)`` will have the length ``3 * len(atoms)//2``. The 3 dimensions of velocity for each odd-indexed atom. New built-in observables may be written by implementing new ``__call__`` methods to return the desired quantities. These should be collected within the ``janus_core.processing.observables`` module. Special cases may also be defined for ease of use, for example:
 
 .. code-block:: python
 

From 874d849d20bff20633267bab79432aceeae5b55d Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Fri, 8 Nov 2024 10:49:10 +0000
Subject: [PATCH 33/37] Add slicelike validator

---
 janus_core/helpers/janus_types.py |  1 +
 janus_core/helpers/utils.py       | 31 +++++++++++++++++++++-
 tests/test_utils.py               | 44 +++++++++++++++++++++++++++++--
 3 files changed, 73 insertions(+), 3 deletions(-)

diff --git a/janus_core/helpers/janus_types.py b/janus_core/helpers/janus_types.py
index 817c5fc2..27c6c556 100644
--- a/janus_core/helpers/janus_types.py
+++ b/janus_core/helpers/janus_types.py
@@ -79,6 +79,7 @@ class PostProcessKwargs(TypedDict, total=False):
     vaf_step: int
     vaf_output_file: PathLike | None
 
+
 class CorrelationKwargs(TypedDict, total=True):
     """Arguments for on-the-fly correlations <ab>."""
 
diff --git a/janus_core/helpers/utils.py b/janus_core/helpers/utils.py
index b0a93e4c..b3b8bd98 100644
--- a/janus_core/helpers/utils.py
+++ b/janus_core/helpers/utils.py
@@ -416,6 +416,34 @@ def track_progress(sequence: Sequence | Iterable, description: str) -> Iterable:
         yield from progress.track(sequence, description=description)
 
 
+def validate_slicelike(maybe_slicelike: SliceLike):
+    """
+    Raise an exception if slc is not a valid SliceLike.
+
+    Parameters
+    ----------
+    maybe_slicelike : SliceLike
+        Candidate to test.
+
+    Raises
+    ------
+    ValueError
+        If maybe_slicelike is not SliceLike.
+    """
+    if isinstance(maybe_slicelike, (slice, range, int)):
+        return
+    if isinstance(maybe_slicelike, tuple) and len(maybe_slicelike) == 3:
+        start, stop, step = maybe_slicelike
+        if (
+            (start is None or isinstance(start, int))
+            and (stop is None or isinstance(stop, int))
+            and isinstance(step, int)
+        ):
+            return
+
+    raise ValueError(f"{maybe_slicelike} is not a valid SliceLike")
+
+
 def slicelike_to_startstopstep(index: SliceLike) -> StartStopStep:
     """
     Standarize `SliceLike`s into tuple of `start`, `stop`, `step`.
@@ -430,6 +458,7 @@ def slicelike_to_startstopstep(index: SliceLike) -> StartStopStep:
     StartStopStep
         Standardized `SliceLike` as `start`, `stop`, `step` triplet.
     """
+    validate_slicelike(index)
     if isinstance(index, int):
         if index == -1:
             return (index, None, 1)
@@ -441,7 +470,7 @@ def slicelike_to_startstopstep(index: SliceLike) -> StartStopStep:
     return index
 
 
-def selector_len(slc: Union[SliceLike, list], selectable_length: int) -> int:
+def selector_len(slc: SliceLike | list, selectable_length: int) -> int:
     """
     Calculate the length of a selector applied to an indexable of a given length.
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 4e9baa11..4423ba5e 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Union
 
 from ase import Atoms
 from ase.io import read
@@ -17,6 +16,7 @@
     none_to_dict,
     selector_len,
     slicelike_to_startstopstep,
+    validate_slicelike,
 )
 
 DATA_PATH = Path(__file__).parent / "data/NaCl.cif"
@@ -198,6 +198,46 @@ def test_slicelike_to_startstopstep(slc: SliceLike, expected: StartStopStep):
         ([0, -1, 2, 9], 10, 4),
     ],
 )
-def test_selector_len(slc: Union[SliceLike, list[int]], len: int, expected: int):
+def test_selector_len(slc: SliceLike | list[int], len: int, expected: int):
     """Test converting SliceLike to StartStopStep."""
     assert selector_len(slc, len) == expected
+
+
+@pytest.mark.parametrize(
+    "slc",
+    [
+        slice(0, 1, 1),
+        slice(0, None, 1),
+        range(3),
+        range(0, 10, 1),
+        -1,
+        0,
+        1,
+        (0),
+        (0, 1, 1),
+        (0, None, 1),
+    ],
+)
+def test_valid_slicelikes(slc):
+    """Test validate_slicelike on valid SliceLikes."""
+    validate_slicelike(slc)
+
+
+@pytest.mark.parametrize(
+    "slc",
+    [
+        1.0,
+        "",
+        None,
+        [1],
+        (None, 0, None),
+        (0, 1, None),
+        (None, None, None),
+        (0, 1),
+        (0, 1, 2, 3),
+    ],
+)
+def test_invalid_slicelikes(slc):
+    """Test validate_slicelike on invalid SliceLikes."""
+    with pytest.raises(ValueError):
+        validate_slicelike(slc)

From 098a4daf1486cc36dc1880f654de8f09aa7bcde1 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <33522054+harveydevereux@users.noreply.github.com>
Date: Fri, 8 Nov 2024 12:05:20 +0000
Subject: [PATCH 34/37] Apply suggestions from code review

validate slicelike return hint

Co-authored-by: Jacob Wilkins <46597752+oerc0122@users.noreply.github.com>
---
 janus_core/helpers/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/janus_core/helpers/utils.py b/janus_core/helpers/utils.py
index b3b8bd98..b49aeba4 100644
--- a/janus_core/helpers/utils.py
+++ b/janus_core/helpers/utils.py
@@ -416,7 +416,7 @@ def track_progress(sequence: Sequence | Iterable, description: str) -> Iterable:
         yield from progress.track(sequence, description=description)
 
 
-def validate_slicelike(maybe_slicelike: SliceLike):
+def validate_slicelike(maybe_slicelike: SliceLike) -> None:
     """
     Raise an exception if slc is not a valid SliceLike.
 

From 93f75efe456492343a4f9e3123a9c8125ebed666 Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Fri, 8 Nov 2024 14:53:41 +0000
Subject: [PATCH 35/37] Remove value_count, update dev guide

---
 docs/source/developer_guide/tutorial.rst | 71 +++++++++++++++++-------
 janus_core/processing/correlator.py      | 26 ++++-----
 janus_core/processing/observables.py     | 60 ++------------------
 3 files changed, 70 insertions(+), 87 deletions(-)

diff --git a/docs/source/developer_guide/tutorial.rst b/docs/source/developer_guide/tutorial.rst
index 76f392fc..31ba7ecf 100644
--- a/docs/source/developer_guide/tutorial.rst
+++ b/docs/source/developer_guide/tutorial.rst
@@ -187,54 +187,87 @@ Alternatively, using ``tox``::
 Adding a new Observable
 =======================
 
-Additional built-in observable quantities may be added for use by the ``janus_core.processing.correlator.Correlation`` class. These should extend ``janus_core.processing.observables.Observable``. The abstract method ``__call__`` should be implemented to obtain the values of the observed quantity from an ``Atoms`` object, and potentially ``value_count`` may need to be overloaded indicating the number of values returned by ``__call__``. The number of values returned will depend on the dimension of the underlying ``Observable``, and potentially the number of atoms resulting from the ``Observable``'s slice. These values are correlated individually and averaged.
+A ``janus_core.processing.observables.Observable`` abstracts obtaining a quantity derived from ``Atoms``. They may be used as kernels for input into analysis such as a correlation.
 
-As an example of building a new ``Observable`` consider the ``janus_core.processing.observables.Stress`` and ``janus_core.processing.observables.Velocity`` built-ins. ``janus_core.processing.observables.Stress`` returns the required stress components as calculated by ``Atoms.get_stress`` called upon the tracked atoms. Therefore ``__call__`` returns one ``float`` per component correlated. The call method applies the atom slice and returns the desired stress components.
+Additional built-in observable quantities may be added for use by the ``janus_core.processing.correlator.Correlation`` class. These should extend ``janus_core.processing.observables.Observable`` and are implemented within the ``janus_core.processing.observables`` module.
+
+The abstract method ``__call__`` should be implemented to obtain the values of the observed quantity from an ``Atoms`` object. When used as part of a ``janus_core.processing.correlator.Correlation``, each value will be correlated and the results averaged.
+
+As an example of building a new ``Observable`` consider the ``janus_core.processing.observables.Stress`` built-in. The following steps may be taken:
+
+1. Defining the observable.
+---------------------------
+
+The stress tensor may be computed on an atoms object using ``Atoms.get_stress``. A user may wish to obtain a particular component, or perhaps only compute the stress on some subset of ``Atoms``. For example during a ``janus_core.calculations.md.MolecularDynamics`` run a user may wish to correlate only the off-diagonal components (shear stress), computed across all atoms.
+
+2. Writing the ``__call__`` method.
+-----------------------------------
+
+In the call method we can use the base ``janus_core.processing.observables.Observable``'s optional atom selector ``atoms_slice`` to first define the subset of atoms to compute the stress for:
 
 .. code-block:: python
 
     def __call__(self, atoms: Atoms) -> list[float]:
         sliced_atoms = atoms[self.atoms_slice]
+        # must be re-attached after slicing for get_stress
         sliced_atoms.calc = atoms.calc
-        stresses = (
+
+Next the stresses may be obtained from:
+
+.. code-block:: python
+
+    stresses = (
             sliced_atoms.get_stress(
                 include_ideal_gas=self.include_ideal_gas, voigt=True
             )
             / units.GPa
         )
-        return stresses[self._indices]
 
-For example an observable representing the ``xy`` and ``zy`` components of stress computed across all odd-index atoms in some ``Atoms`` object can be constructed as follows:
+Finally, to facilitate handling components in a symbolic way, ``janus_core.processing.observables.ComponentMixin`` exists to parse ``str`` symbolic components to ``int`` indices by defining a suitable mapping. For the stress tensor (and the format of ``Atoms.get_stress``) a suitable mapping is defined in ``janus_core.processing.observables.Stress``'s ``__init__`` method:
 
 .. code-block:: python
 
-    s = Stress(components=["xy", "zy"], atoms_slice=(0, None, 2))
-
-In this case ``s(atoms)`` will return 2 values, no matter how many atoms are sliced from ``atoms_slice``.
+        ComponentMixin.__init__(
+            self,
+            components={
+                "xx": 0,
+                "yy": 1,
+                "zz": 2,
+                "yz": 3,
+                "zy": 3,
+                "xz": 4,
+                "zx": 4,
+                "xy": 5,
+                "yx": 5,
+            },
+        )
 
-However, the ``janus_core.processing.observables.Velocity`` built-in's ``__call__`` not only returns atom velocity for the requested dimensions, but also returns them for every tracked atom, i.e:
+This then concludes the ``__call__`` method for ``janus_core.processing.observables.Stress`` by using ``janus_core.processing.observables.ComponentMixin``'s
+pre-calculated indices:
 
 .. code-block:: python
 
-       def __call__(self, atoms: Atoms) -> list[float]:
-           return atoms.get_velocities()[self.atoms_slice, :][:, self._indices].flatten()
+    return stesses[self._indices]
 
-so ``value_count`` is overloaded to reflect this.
+The combination of the above means a user may obtain, say, the ``xy`` and ``zy`` stress tensor components over odd-indexed atoms by calling the following observable on an ``Atoms``:
 
 .. code-block:: python
 
-       def value_count(self, n_atoms: int | None = None) -> int:
-           return selector_len(self.atoms_slice, n_atoms) * self.dimension
+    s = Stress(components=["xy", "zy"], atoms_slice=(0, None, 2))
+
 
-Unlike ``janus_core.processing.observables.Stress`` when we create a ``janus_core.processing.observables.Velocity`` observable like so,
+Since usually total system stresses are required we can define two built-ins to handle the shear and hydrostatic stresses like so:
 
 .. code-block:: python
 
-    v = Velocity(components=["x", "y", "z"], atoms_slice=(0, None, 2))
+    StressHydrostatic = Stress(components=["xx", "yy", "zz"])
+    StressShear = Stress(components=["xy", "yz", "zx"])
+
+Where by default ``janus_core.processing.observables.Observable``'s ``atoms_slice`` is ``slice(0, None, 1)``, which expands to all atoms in an ``Atoms``.
 
-the value of ``v(atoms)`` will have the length ``3 * len(atoms)//2``. The 3 dimensions of velocity for each odd-indexed atom. New built-in observables may be written by implementing new ``__call__`` methods to return the desired quantities. These should be collected within the ``janus_core.processing.observables`` module. Special cases may also be defined for ease of use, for example:
+For comparison the ``janus_core.processing.observables.Velocity`` built-in's ``__call__`` not only returns atom velocity for the requested components, but also returns them for every tracked atom i.e:
 
 .. code-block:: python
 
-    StressHydrostatic = Stress(components=["xx", "yy", "zz"])
-    StressShear = Stress(components=["xy", "yz", "zx"])
+    def __call__(self, atoms: Atoms) -> list[float]:
+        return atoms.get_velocities()[self.atoms_slice, :][:, self._indices].flatten()
diff --git a/janus_core/processing/correlator.py b/janus_core/processing/correlator.py
index d612cdeb..9f7da0dd 100644
--- a/janus_core/processing/correlator.py
+++ b/janus_core/processing/correlator.py
@@ -232,20 +232,13 @@ def __init__(
             Frequency to update the correlation, md steps.
         """
         self.name = name
+        self.blocks = blocks
+        self.points = points
+        self.averaging = averaging
         self._get_a = a
         self._get_b = b
 
-        a_values = self._get_a.value_count(n_atoms)
-        b_values = self._get_b.value_count(n_atoms)
-
-        if a_values != b_values:
-            raise ValueError("Observables have inconsistent sizes")
-        self._values = a_values
-
-        self._correlators = [
-            Correlator(blocks=blocks, points=points, averaging=averaging)
-            for _ in range(self._values)
-        ]
+        self._correlators = None
         self._update_frequency = update_frequency
 
     @property
@@ -269,8 +262,15 @@ def update(self, atoms: Atoms) -> None:
         atoms : Atoms
             Atoms object to observe values from.
         """
-        atom_pairs = zip(self._get_a(atoms), self._get_b(atoms))
-        for corr, values in zip(self._correlators, atom_pairs):
+        value_pairs = zip(self._get_a(atoms), self._get_b(atoms))
+        if self._correlators is None:
+            self._correlators = [
+                Correlator(
+                    blocks=self.blocks, points=self.points, averaging=self.averaging
+                )
+                for _ in range(len(self._get_a(atoms)))
+            ]
+        for corr, values in zip(self._correlators, value_pairs):
             corr.update(*values)
 
     def get(self) -> tuple[Iterable[float], Iterable[float]]:
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 6969820c..00b1de97 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -10,7 +10,7 @@
 if TYPE_CHECKING:
     from janus_core.helpers.janus_types import SliceLike
 
-from janus_core.helpers.utils import selector_len, slicelike_to_startstopstep
+from janus_core.helpers.utils import slicelike_to_startstopstep
 
 
 # pylint: disable=too-few-public-methods
@@ -22,13 +22,9 @@ class Observable(ABC):
     ----------
     atoms_slice : list[int] | SliceLike | None = None
         A slice of atoms to observe.
-    dimension : int
-        The dimension of the observed data.
     """
 
-    def __init__(
-        self, atoms_slice: list[int] | SliceLike | None = None, dimension: int = 1
-    ):
+    def __init__(self, atoms_slice: list[int] | SliceLike | None = None):
         """
         Initialise an observable with a given dimensionality.
 
@@ -36,10 +32,7 @@ def __init__(
         ----------
         atoms_slice : list[int] | SliceLike | None = None
             A slice of atoms to observe.
-        dimension : int
-            The dimension of the observed data.
         """
-        self._dimension = dimension
         if atoms_slice:
             if isinstance(atoms_slice, list):
                 self.atoms_slice = atoms_slice
@@ -64,34 +57,6 @@ def __call__(self, atoms: Atoms) -> list[float]:
             The observed value, with dimensions atoms by self.dimension.
         """
 
-    def value_count(self, n_atoms: int | None = None) -> int:
-        """
-        Count of values returned by __call__.
-
-        Parameters
-        ----------
-        n_atoms : int | None
-            Atom count to expand atoms_slice.
-
-        Returns
-        -------
-        int
-            The number of values returned by __call__.
-        """
-        return self.dimension
-
-    @property
-    def dimension(self):
-        """
-        Dimension of the observable. Commensurate with self.__call__.
-
-        Returns
-        -------
-        int
-            Observables dimension.
-        """
-        return self._dimension
-
 
 class ComponentMixin:
     """
@@ -199,7 +164,7 @@ def __init__(
         )
         self._set_components(components)
 
-        Observable.__init__(self, atoms_slice, len(components))
+        Observable.__init__(self, atoms_slice)
         self.include_ideal_gas = include_ideal_gas
 
     def __call__(self, atoms: Atoms) -> list[float]:
@@ -217,6 +182,7 @@ def __call__(self, atoms: Atoms) -> list[float]:
             The stress components in GPa units.
         """
         sliced_atoms = atoms[self.atoms_slice]
+        # must be re-attached after slicing for get_stress
         sliced_atoms.calc = atoms.calc
         stresses = (
             sliced_atoms.get_stress(
@@ -263,7 +229,7 @@ def __init__(
         ComponentMixin.__init__(self, components={"x": 0, "y": 1, "z": 2})
         self._set_components(components)
 
-        Observable.__init__(self, atoms_slice, len(components))
+        Observable.__init__(self, atoms_slice)
 
     def __call__(self, atoms: Atoms) -> list[float]:
         """
@@ -280,19 +246,3 @@ def __call__(self, atoms: Atoms) -> list[float]:
             The velocity values.
         """
         return atoms.get_velocities()[self.atoms_slice, :][:, self._indices].flatten()
-
-    def value_count(self, n_atoms: int | None = None) -> int:
-        """
-        Count of values returned by __call__.
-
-        Parameters
-        ----------
-        n_atoms : int | None
-            Atom count to expand atoms_slice.
-
-        Returns
-        -------
-        int
-            The number of values returned by __call__.
-        """
-        return selector_len(self.atoms_slice, n_atoms) * self.dimension

From d1d7e4718d0c4323ff2923f209420ad77ba5410f Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Fri, 8 Nov 2024 15:31:51 +0000
Subject: [PATCH 36/37] Check atoms is Atoms, clearer exception

---
 janus_core/processing/observables.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 00b1de97..2d8d9083 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -180,7 +180,14 @@ def __call__(self, atoms: Atoms) -> list[float]:
         -------
         list[float]
             The stress components in GPa units.
+
+        Raises
+        ------
+        ValueError
+            If atoms is not an Atoms object.
         """
+        if not isinstance(atoms, Atoms):
+            raise ValueError("atoms should be an Atoms object")
         sliced_atoms = atoms[self.atoms_slice]
         # must be re-attached after slicing for get_stress
         sliced_atoms.calc = atoms.calc

From 6c3071223209e1eed67cf80cb444cbb4ec47ecea Mon Sep 17 00:00:00 2001
From: Harvey Devereux <h.devereux@qmul.ac.uk>
Date: Fri, 8 Nov 2024 15:37:28 +0000
Subject: [PATCH 37/37] remove uneeded

---
 docs/source/conf.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/docs/source/conf.py b/docs/source/conf.py
index 51483d12..18589ceb 100755
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -209,6 +209,4 @@
     ("py:class", "Calculator"),
     ("py:class", "Context"),
     ("py:class", "Path"),
-    ("py:obj", "dimension"),
-    ("py:obj", "janus_core.processing.observables.Stress.value_count")
 ]