Add transition_counts() and transition_probabilities() methods

Add methods to aggregate per-site transition data into a TransitionTable,
a labelled square matrix with convenient access patterns (.matrix,
.get(), .to_dict(), .reorder()).

- transition_counts(by='site'|'label') returns raw hop counts
- transition_probabilities(by='site'|'label') returns row-normalised
  probabilities (rows with no outgoing transitions remain as zeros)
- by='label' aggregates across sites sharing a label; unlabelled sites
  are skipped with a warning if transitions are dropped
- Optional keys parameter for custom row/column ordering
- Unknown destination indices raise ValueError

Closes #45

Author: bjmorgan

site_analysis/__init__.py

@@ -19,3 +19,4 @@
     create_trajectory_with_polyhedral_sites,
     create_trajectory_with_dynamic_voronoi_sites
 )
+from site_analysis.transition_table import TransitionTable

site_analysis/trajectory.py

@@ -27,14 +27,18 @@
     ...              .build())
 """
 
+import warnings
 from collections import Counter
 from collections.abc import Iterable
-from typing import Sequence
+from typing import Literal, Sequence
 
+import numpy as np
 from tqdm.auto import tqdm
 
 from pymatgen.core import Structure
 
+from .transition_table import TransitionTable
+
 from .atom import Atom
 from .dynamic_voronoi_site import DynamicVoronoiSite
 from .dynamic_voronoi_site_collection import DynamicVoronoiSiteCollection
@@ -161,7 +165,126 @@ def site_labels(self) -> list[str | None]:
             A list of site labels (or None for sites without labels).
         """
         return [s.label for s in self.sites]
-   
+
+    def transition_counts(
+        self,
+        by: Literal['site', 'label'] = 'site',
+        *,
+        keys: Sequence[int] | Sequence[str] | None = None,
+    ) -> TransitionTable:
+        """Return transition counts as a :class:`TransitionTable`.
+
+        Args:
+            by: Aggregation key. ``'site'`` (default) keys by site index;
+                ``'label'`` aggregates by site label (unlabelled sites are
+                skipped, with a warning if any transitions are dropped).
+            keys: Optional key ordering for rows and columns. If ``None``,
+                keys are sorted. Must be a permutation of the default keys.
+
+        Returns:
+            A :class:`TransitionTable` of integer counts.
+
+        Raises:
+            ValueError: If ``by`` is not ``'site'`` or ``'label'``, if
+                *keys* does not match the default key set, or if a site
+                has a transition to an unknown site index.
+        """
+        if by == 'site':
+            site_keys = tuple(sorted(s.index for s in self.sites))
+            index_set = set(site_keys)
+            idx_lookup = {k: i for i, k in enumerate(site_keys)}
+            n = len(site_keys)
+            matrix = np.zeros((n, n), dtype=int)
+            for site in self.sites:
+                row = idx_lookup[site.index]
+                for dest, count in site.transitions.items():
+                    if dest not in index_set:
+                        raise ValueError(
+                            f"Site {site.index} has a transition to unknown "
+                            f"site index {dest}."
+                        )
+                    matrix[row, idx_lookup[dest]] = count
+            result_keys: tuple[int, ...] | tuple[str, ...] = site_keys
+        elif by == 'label':
+            all_site_indices = {s.index for s in self.sites}
+            index_to_label = {
+                s.index: s.label for s in self.sites if s.label is not None
+            }
+            label_keys = tuple(sorted(set(index_to_label.values())))
+            label_lookup = {k: i for i, k in enumerate(label_keys)}
+            n = len(label_keys)
+            matrix = np.zeros((n, n), dtype=int)
+            dropped = 0
+            for site in self.sites:
+                src_label = index_to_label.get(site.index)
+                if src_label is None:
+                    for dest in site.transitions:
+                        if dest not in all_site_indices:
+                            raise ValueError(
+                                f"Site {site.index} has a transition to unknown "
+                                f"site index {dest}."
+                            )
+                    dropped += sum(site.transitions.values())
+                    continue
+                for dest, count in site.transitions.items():
+                    if dest not in all_site_indices:
+                        raise ValueError(
+                            f"Site {site.index} has a transition to unknown "
+                            f"site index {dest}."
+                        )
+                    dest_label = index_to_label.get(dest)
+                    if dest_label is None:
+                        dropped += count
+                        continue
+                    matrix[label_lookup[src_label], label_lookup[dest_label]] += count
+            if dropped > 0:
+                warnings.warn(
+                    f"{dropped} transition(s) involving unlabelled sites were "
+                    f"excluded from the label-aggregated counts.",
+                    stacklevel=2,
+                )
+            result_keys = label_keys
+        else:
+            raise ValueError(f"Invalid value for 'by': {by!r}. Must be 'site' or 'label'.")
+        table = TransitionTable(keys=result_keys, matrix=matrix)
+        if keys is not None:
+            return table.reorder(keys)
+        return table
+
+    def transition_probabilities(
+        self,
+        by: Literal['site', 'label'] = 'site',
+        *,
+        keys: Sequence[int] | Sequence[str] | None = None,
+    ) -> TransitionTable:
+        """Return row-normalised transition probabilities as a :class:`TransitionTable`.
+
+        Each row is normalised so that its values sum to 1.0. Rows with no
+        outgoing transitions remain as all zeros.
+
+        Args:
+            by: Aggregation key. ``'site'`` (default) keys by site index;
+                ``'label'`` aggregates by site label (unlabelled sites are
+                skipped, with a warning if any transitions are dropped).
+            keys: Optional key ordering for rows and columns. If ``None``,
+                keys are sorted. Must be a permutation of the default keys.
+
+        Returns:
+            A :class:`TransitionTable` of float probabilities.
+
+        Raises:
+            ValueError: If ``by`` is not ``'site'`` or ``'label'``, if
+                *keys* does not match the default key set, or if a site
+                has a transition to an unknown site index.
+        """
+        counts = self.transition_counts(by=by, keys=keys)
+        count_data = counts.matrix.astype(float)
+        row_sums = count_data.sum(axis=1)
+        probs = np.zeros_like(count_data)
+        nonzero = row_sums > 0
+        probs[nonzero] = count_data[nonzero] / row_sums[nonzero, np.newaxis]
+        return TransitionTable(keys=counts.keys, matrix=probs)
+
     @property
     def atom_sites(self) -> list[int | None]:
         """Return the sites that each atom currently occupies.

site_analysis/transition_table.py

@@ -0,0 +1,158 @@
+"""Transition table for storing labelled transition data.
+
+Provides the :class:`TransitionTable` class, which stores transition
+counts or probabilities as a labelled square matrix with convenient
+access patterns.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence, overload
+
+import numpy as np
+
+
+class TransitionTable:
+    """A labelled square matrix of transition data.
+
+    Stores transition counts or probabilities with named keys for
+    rows and columns. Provides multiple access patterns:
+
+    - ``.matrix`` — the raw (read-only) :class:`numpy.ndarray`
+    - ``.get(from_key, to_key)`` — key-based lookup (matched types enforced)
+    - ``.to_dict()`` — square dict-of-dicts
+    - ``.reorder(keys)`` — return a new table with reordered rows/columns
+
+    Args:
+        keys: Row and column labels (site indices or site labels).
+        matrix: A square 2-D numpy array of transition values.
+
+    Raises:
+        ValueError: If *matrix* is not 2-D and square, if
+            ``len(keys) != matrix.shape[0]``, or if *keys* contains
+            duplicates.
+    """
+
+    __slots__ = ('_keys', '_matrix', '_key_to_index', '_frozen')
+
+    def __init__(
+        self,
+        keys: tuple[int, ...] | tuple[str, ...],
+        matrix: np.ndarray,
+    ) -> None:
+        self._matrix = np.array(matrix, copy=True)
+        if self._matrix.ndim != 2 or self._matrix.shape[0] != self._matrix.shape[1]:
+            raise ValueError(
+                f"matrix must be a square 2-D array, got shape {self._matrix.shape}"
+            )
+        if len(keys) != self._matrix.shape[0]:
+            raise ValueError(
+                f"len(keys) ({len(keys)}) != matrix dimension "
+                f"({self._matrix.shape[0]})"
+            )
+        if len(set(keys)) != len(keys):
+            raise ValueError("keys must not contain duplicates")
+        self._keys = keys
+        self._key_to_index = {k: i for i, k in enumerate(keys)}
+        self._matrix.flags.writeable = False
+        self._frozen = True
+
+    @property
+    def keys(self) -> tuple[int, ...] | tuple[str, ...]:
+        """Row and column labels."""
+        return self._keys
+
+    @property
+    def matrix(self) -> np.ndarray:
+        """The transition data as a read-only 2-D numpy array."""
+        return self._matrix
+
+    @overload
+    def get(self, from_key: int, to_key: int) -> int | float: ...
+    @overload
+    def get(self, from_key: str, to_key: str) -> int | float: ...
+
+    def get(self, from_key: int | str, to_key: int | str) -> int | float:
+        """Look up a single transition value by key.
+
+        Args:
+            from_key: The source key (row).
+            to_key: The destination key (column).
+
+        Returns:
+            The transition value at ``(from_key, to_key)``.
+
+        Raises:
+            KeyError: If either key is not present in the table.
+        """
+        try:
+            i = self._key_to_index[from_key]
+        except KeyError:
+            raise KeyError(from_key) from None
+        try:
+            j = self._key_to_index[to_key]
+        except KeyError:
+            raise KeyError(to_key) from None
+        value: int | float = self._matrix[i, j].item()
+        return value
+
+    def to_dict(self) -> dict[int | str, dict[int | str, int | float]]:
+        """Convert to a square dict-of-dicts.
+
+        Returns:
+            A dict ``{from_key: {to_key: value}}`` mirroring the matrix.
+        """
+        return {
+            self._keys[i]: {
+                self._keys[j]: self._matrix[i, j].item()
+                for j in range(len(self._keys))
+            }
+            for i in range(len(self._keys))
+        }
+
+    def reorder(self, keys: Sequence[int] | Sequence[str]) -> TransitionTable:
+        """Return a new table with rows and columns reordered.
+
+        Args:
+            keys: The desired key ordering. Must contain exactly
+                the same keys as the current table.
+
+        Returns:
+            A new :class:`TransitionTable` with reordered rows and columns.
+
+        Raises:
+            ValueError: If *keys* does not match the current key set exactly.
+        """
+        new_keys: tuple[int, ...] | tuple[str, ...] = tuple(keys)  # type: ignore[assignment]
+        if len(new_keys) != len(self._keys) or set(new_keys) != set(self._keys):
+            missing = sorted(set(self._keys) - set(new_keys))
+            extra = sorted(set(new_keys) - set(self._keys))
+            parts = []
+            if missing:
+                parts.append(f"missing keys: {missing!r}")
+            if extra:
+                parts.append(f"unknown keys: {extra!r}")
+            raise ValueError(
+                f"keys must be a permutation of the current keys; {'; '.join(parts)}"
+            )
+        order = [self._key_to_index[k] for k in new_keys]
+        reordered = self._matrix[np.ix_(order, order)]
+        return TransitionTable(keys=new_keys, matrix=reordered)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, TransitionTable):
+            return NotImplemented
+        return self._keys == other._keys and np.array_equal(self._matrix, other._matrix)
+
+    __hash__ = None  # type: ignore[assignment]
+
+    def __repr__(self) -> str:
+        return (
+            f"TransitionTable(keys={self._keys!r}, "
+            f"shape={self._matrix.shape[0]}x{self._matrix.shape[1]})"
+        )
+
+    def __setattr__(self, name: str, value: object) -> None:
+        if getattr(self, '_frozen', False):
+            raise AttributeError("TransitionTable is immutable")
+        super().__setattr__(name, value)