Clarify values

Use kwargs only in inits

atoms -> atoms_slice

replace atom_count method with value_count method

Stress now can slice atoms

Author: harveydevereux

janus_core/helpers/correlator.py

@@ -243,18 +243,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
@@ -298,11 +298,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:

janus_core/helpers/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,37 +52,33 @@ def __call__(self, atoms: Atoms, *args, **kwargs) -> list[float]:
             The observed value, with dimensions atoms by self.dimension.
         """
 
-    @property
-    def dimension(self):
+    def value_count(self, n_atoms: int | None = None) -> int:
         """
-        Dimension of the observable. Commensurate with self.__call__.
+        Count of values returned by __call__.
+
+        Parameters
+        ----------
+        n_atoms : int | None
+            Atom count to expand atoms_slice.
 
         Returns
         -------
         int
-            Observables dimension.
+            The number of values returned by __call__.
         """
-        return self._dimension
+        return self.dimension
 
-    def atom_count(self, n_atoms: int):
+    @property
+    def dimension(self):
         """
-        Atom count to average over.
-
-        Parameters
-        ----------
-        n_atoms : int
-            Total possible atoms.
+        Dimension of the observable. Commensurate with self.__call__.
 
         Returns
         -------
         int
-            Atom count averaged over.
+            Observables dimension.
         """
-        if self.atoms:
-            if isinstance(self.atoms, list):
-                return len(self.atoms)
-            return slicelike_len_for(self.n_atoms)
-        return 0
+        return self._dimension
 
 
 class ComponentMixin:
@@ -163,18 +158,28 @@ 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.
 
         Parameters
         ----------
         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.
         """
@@ -194,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
 
@@ -215,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
@@ -234,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.
@@ -250,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]:
         """
@@ -280,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

tests/test_correlator.py

@@ -94,17 +94,17 @@ 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,
                 "averaging": 1,
                 "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,
@@ -128,8 +128,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):
@@ -154,8 +154,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,