Add docstrings for the vector_mapping and universe_mapping modules.

Add real benchmarking tests for vector mapping operations.

Author: jathoms

justfile

@@ -32,3 +32,7 @@ docs open="" output_dir="docs/_build/html":
 doctest:
   uv sync --group docs
   uv run sphinx-build -b doctest -W --keep-going docs/source docs/_build/doctest
+
+benchmark:
+  uv sync --group dev
+  uv run pytest tests/benchmark -m benchmark

pyproject.toml

@@ -27,6 +27,7 @@ dev = [
   "pyrefly>=0.29.2",
   "pyright>=1.1.404",
   "pytest>=8.3.0",
+  "pytest-benchmark>=5.2.3",
   "pytest-cov>=5.0.0",
   "ruff>=0.12.8",
   "ty>=0.0.1a17",
@@ -75,6 +76,9 @@ ignore = ["EM101"]
 [tool.ruff.format]
 docstring-code-format = true
 
+[tool.pytest.ini_options]
+addopts = "-m 'not benchmark'"
+
 [tool.ty.src]
 include = ["src"]
 exclude = ["src/backtest_lib/examples/*"]

src/backtest_lib/universe/universe_mapping.py

@@ -1,11 +1,17 @@
+"""Universe-aligned numeric mappings.
+
+Universe mappings are vector mappings keyed by security identifiers. Backends
+can use aligned key ordering to perform fast vectorized arithmetic. When two
+mappings share the same universe ordering, operations are fastest; mismatched
+ordering remains correct but incurs extra reindexing. See
+``tests/benchmark/polars_impl/test_universe_mapping_benchmark.py`` for benchmarks.
+"""
+
 from __future__ import annotations
 
-from abc import ABC
+from abc import ABC, abstractmethod
 from collections.abc import Iterable, Mapping
-from typing import (
-    TypeVar,
-    abstractmethod,
-)
+from typing import TypeVar
 
 from backtest_lib.market._backends import _get_mapping_type_from_backend
 from backtest_lib.market.plotting import UniverseMappingPlotAccessor
@@ -16,33 +22,63 @@
 
 
 class UniverseMapping[Scalar: (float, int)](VectorMapping[str, Scalar], ABC):
+    """Vector mapping keyed by security identifiers.
+
+    Implementations back this interface with dense vectors aligned to a universe
+    ordering, enabling fast arithmetic when operands share the same ordering.
+    """
+
     @property
-    def plot(self) -> UniverseMappingPlotAccessor: ...
+    def plot(self) -> UniverseMappingPlotAccessor:
+        """Return the plotting accessor for the mapping."""
+        ...
 
     @abstractmethod
     def __truediv__(
         self,
         other: VectorMapping | float | int | Mapping[str, int | float],
-    ) -> UniverseMapping[float]: ...
+    ) -> UniverseMapping[float]:
+        """Return the elementwise quotient of two mappings."""
+        ...
 
     @abstractmethod
     def __rtruediv__(
         self,
         other: VectorMapping | float | int | Mapping[str, int | float],
-    ) -> UniverseMapping[float]: ...
+    ) -> UniverseMapping[float]:
+        """Return the elementwise quotient with reversed operands."""
+        ...
 
     @abstractmethod
-    def floor(self) -> UniverseMapping[int]: ...
+    def floor(self) -> UniverseMapping[int]:
+        """Return a mapping with values floored to integers."""
+        ...
 
     @abstractmethod
-    def truncate(self) -> UniverseMapping[int]: ...
+    def truncate(self) -> UniverseMapping[int]:
+        """Return a mapping with values truncated to integers."""
+        ...
 
 
 def make_universe_mapping[T: (int, float)](
     m: Mapping[str, T],
     universe: Iterable[str],
     constructor_backend: str = "polars",
 ) -> UniverseMapping[T]:
+    """Create a universe-aligned mapping from an arbitrary mapping.
+
+    Values are re-ordered to match ``universe`` and missing keys are filled with
+    zeros. Keeping a stable universe ordering enables faster vectorized
+    operations when combining mappings.
+
+    Args:
+        m: Mapping of security identifiers to values.
+        universe: Ordered universe defining the mapping alignment.
+        constructor_backend: Backend used for the concrete mapping type.
+
+    Returns:
+        UniverseMapping aligned to ``universe``.
+    """
     if not isinstance(m, UniverseMapping) and isinstance(m, Mapping):
         backend_mapping_type = _get_mapping_type_from_backend(constructor_backend)
         # TODO: check if this dictionary collection is slow.

src/backtest_lib/universe/vector_mapping.py

@@ -1,3 +1,13 @@
+"""Vectorized mapping primitives for numeric data.
+
+These mappings support elementwise arithmetic and can be implemented with
+backend-specific vectorized operations. Performance is best when two mappings
+share the same key order, because backends can apply operations without
+reordering. The benchmarks in
+``tests/benchmark/polars_impl/test_universe_mapping_benchmark.py`` illustrate the
+penalty for mismatched ordering.
+"""
+
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
@@ -12,6 +22,13 @@
 
 
 class VectorMapping[K, V_co](Mapping[K, V_co], ABC):
+    """Mapping with vectorized arithmetic semantics.
+
+    Backends may rely on matching key order to perform fast vector operations.
+    When key order differs between two mappings, operations remain correct but
+    typically require reindexing and incur a performance cost.
+    """
+
     @overload
     def __add__(
         self: VectorMapping[K, int],
@@ -32,7 +49,9 @@ def __add__(
         | Mapping[K, int | float],
     ) -> VectorMapping[K, float]: ...
     @abstractmethod
-    def __add__(self, other) -> VectorMapping: ...
+    def __add__(self, other) -> VectorMapping:
+        """Return the elementwise sum of two mappings."""
+        ...
 
     @overload
     def __radd__(
@@ -54,7 +73,9 @@ def __radd__(
         | Mapping[K, int | float],
     ) -> VectorMapping[K, float]: ...
     @abstractmethod
-    def __radd__(self, other) -> VectorMapping: ...
+    def __radd__(self, other) -> VectorMapping:
+        """Return the elementwise sum with reversed operands."""
+        ...
 
     @overload
     def __sub__(
@@ -76,7 +97,9 @@ def __sub__(
         | Mapping[K, int | float],
     ) -> VectorMapping[K, float]: ...
     @abstractmethod
-    def __sub__(self, other) -> VectorMapping: ...
+    def __sub__(self, other) -> VectorMapping:
+        """Return the elementwise difference of two mappings."""
+        ...
 
     @overload
     def __rsub__(
@@ -98,7 +121,9 @@ def __rsub__(
         | Mapping[K, int | float],
     ) -> VectorMapping[K, float]: ...
     @abstractmethod
-    def __rsub__(self, other) -> VectorMapping: ...
+    def __rsub__(self, other) -> VectorMapping:
+        """Return the elementwise difference with reversed operands."""
+        ...
 
     @overload
     def __mul__(
@@ -120,7 +145,9 @@ def __mul__(
         | Mapping[K, int | float],
     ) -> VectorMapping[K, float]: ...
     @abstractmethod
-    def __mul__(self, other) -> VectorMapping: ...
+    def __mul__(self, other) -> VectorMapping:
+        """Return the elementwise product of two mappings."""
+        ...
 
     @overload
     def __rmul__(
@@ -142,21 +169,29 @@ def __rmul__(
         | Mapping[K, int | float],
     ) -> VectorMapping[K, float]: ...
     @abstractmethod
-    def __rmul__(self, other) -> VectorMapping: ...
+    def __rmul__(self, other) -> VectorMapping:
+        """Return the elementwise product with reversed operands."""
+        ...
 
     @abstractmethod
     def __truediv__(
         self,
         other: VectorMapping | float | int | Mapping[K, int | float],
-    ) -> VectorMapping[K, float]: ...
+    ) -> VectorMapping[K, float]:
+        """Return the elementwise quotient of two mappings."""
+        ...
 
     @abstractmethod
     def __rtruediv__(
         self, other: VectorMapping | float | int | Mapping[K, int | float]
-    ) -> VectorMapping[K, float]: ...
+    ) -> VectorMapping[K, float]:
+        """Return the elementwise quotient with reversed operands."""
+        ...
 
     @abstractmethod
-    def sum(self) -> float: ...
+    def sum(self) -> float:
+        """Return the sum of all values."""
+        ...
 
     # Default implementations of mean(), override this
     # where possible with faster vector operations
@@ -167,20 +202,32 @@ def mean(self) -> float:
         return self.sum() / n
 
     @abstractmethod
-    def abs(self) -> Self: ...
+    def abs(self) -> Self:
+        """Return a mapping with absolute values."""
+        ...
 
     @abstractmethod
-    def truncate(self) -> VectorMapping[K, int]: ...
+    def truncate(self) -> VectorMapping[K, int]:
+        """Return a mapping with values truncated to integers."""
+        ...
 
     @abstractmethod
-    def floor(self) -> VectorMapping[K, int]: ...
+    def floor(self) -> VectorMapping[K, int]:
+        """Return a mapping with values floored to integers."""
+        ...
 
     @abstractmethod
-    def __getitem__(self, key: K) -> V_co: ...
+    def __getitem__(self, key: K) -> V_co:
+        """Return the value for a key."""
+        ...
 
     @abstractmethod
-    def __iter__(self) -> Iterator[K]: ...
+    def __iter__(self) -> Iterator[K]:
+        """Return an iterator over keys in order."""
+        ...
 
     @classmethod
     @abstractmethod
-    def from_vectors(cls, keys: Iterable[K_contra], values: Iterable[V_co]) -> Self: ...
+    def from_vectors(cls, keys: Iterable[K_contra], values: Iterable[V_co]) -> Self:
+        """Create a mapping from ordered key/value vectors."""
+        ...

tests/benchmark/polars_impl/test_universe_mapping_benchmark.py

@@ -0,0 +1,47 @@
+"""Benchmarks for UniverseMapping ordering behavior.
+
+These timings highlight the performance impact when key order differs between mappings.
+"""
+
+from __future__ import annotations
+
+import pytest
+from pytest_benchmark.fixture import BenchmarkFixture
+
+from backtest_lib.market.polars_impl import PolarsUniverseMapping
+
+
+def _benchmark_addition(
+    keys: list[str],
+    other_keys: list[str],
+) -> PolarsUniverseMapping:
+    values = [1] * len(keys)
+    acc = PolarsUniverseMapping.from_vectors(keys, values)
+    other = PolarsUniverseMapping.from_vectors(other_keys, values)
+    return acc + other
+
+
+@pytest.mark.benchmark
+def test_small_ordering_same(benchmark: BenchmarkFixture) -> None:
+    keys = ["a", "b", "c"]
+    benchmark(_benchmark_addition, keys, keys)
+
+
+@pytest.mark.benchmark
+def test_small_ordering_diff(benchmark: BenchmarkFixture) -> None:
+    keys = ["a", "b", "c"]
+    diff_keys = ["c", "a", "b"]
+    benchmark(_benchmark_addition, keys, diff_keys)
+
+
+@pytest.mark.benchmark
+def test_large_ordering_same(benchmark: BenchmarkFixture) -> None:
+    keys = [str(i) for i in range(1000)]
+    benchmark(_benchmark_addition, keys, keys)
+
+
+@pytest.mark.benchmark
+def test_large_ordering_diff(benchmark: BenchmarkFixture) -> None:
+    keys = [str(i) for i in range(1000)]
+    diff_keys = list(reversed(keys))
+    benchmark(_benchmark_addition, keys, diff_keys)

tests/conftest.py

@@ -6,6 +6,13 @@
 from backtest_lib import MarketView
 
 
+def pytest_configure(config: pytest.Config) -> None:
+    config.addinivalue_line(
+        "markers",
+        "benchmark: performance benchmarks (opt-in)",
+    )
+
+
 @pytest.fixture(scope="session")
 def test_data_dir() -> Path:
     return Path(__file__).resolve().parent / "data"