Metrics¶

def update(self, values: Union[np.ndarray, Iterable[float], float]):
    """Update aggregates with a batch of scalar values.

    The input is converted to ``np.float32`` and flattened. If the resulting
    array is empty, the call is a no-op.

    Args:
        values: A scalar or array-like collection of numeric values.
    """
    v = np.asarray(values, dtype=np.float32).ravel()
    m = int(v.size)
    if m == 0:
        return

    if self.stats is None:
        self.n = m
        self.stats = {
            "mean": float(np.mean(v)),
            "mean_squares": float(np.mean(v**2)),
            "max": float(np.max(v)),
            "min": float(np.min(v)),
            "mag": float(np.max(np.abs(v))),
        }
        return

    # Weighted update of first and second moments.
    n_prev = self.n
    self.n += m
    w_new = m / self.n
    w_old = (self.n - m) / self.n

    self.stats = {
        "mean": float(np.mean(v)) * w_new + float(self.stats["mean"]) * w_old,
        "mean_squares": float(np.mean(v**2)) * w_new
        + float(self.stats["mean_squares"]) * w_old,
        "max": float(max(float(np.max(v)), float(self.stats["max"]))),
        "min": float(min(float(np.min(v)), float(self.stats["min"]))),
        "mag": float(max(float(np.max(np.abs(v))), float(self.stats["mag"]))),
    }

def get(self) -> Dict[str, float]:
    r"""Return the current statistics as a flat dict of scalars.

    Returns:
        A mapping containing:

        - ``mean``: :math:`\mu`
        - ``std``: :math:`\sigma = \sqrt{\max(0, \mathbb{E}[X^2] - \mu^2)}`
        - ``max``: :math:`\max x`
        - ``min``: :math:`\min x`
        - ``mag``: :math:`\max |x|`

        If :attr:`prefix` is non-empty, keys are prefixed with
        ``"{prefix}_"``.

    Raises:
        RuntimeError: If called before any data has been observed.
    """
    if self.stats is None:
        raise RuntimeError("Stats.get() called before any update().")

    mean = float(self.stats["mean"])
    mean_sq = float(self.stats["mean_squares"])
    std = float(np.sqrt(np.maximum(0.0, mean_sq - mean**2)))

    out: Dict[str, float] = {
        "mean": mean,
        "std": std,
        "max": float(self.stats["max"]),
        "min": float(self.stats["min"]),
        "mag": float(self.stats["mag"]),
    }
    if self.prefix:
        out = {f"{self.prefix}_{k}": v for k, v in out.items()}
    return out

def __add__(self, other: "Stats") -> "Stats":
    """Combine two :class:`Stats` objects into a new aggregate.

    This is useful when you have two independent streams and want the same
    summary you would have obtained if you had processed all samples in one
    stream. The combination is exact for the tracked aggregates
    (first/second moments and extrema).

    Args:
        other: Another :class:`Stats` instance with the same :attr:`prefix`.

    Returns:
        A new :class:`Stats` instance representing the combined aggregate.

    Raises:
        AssertionError: If :attr:`prefix` differs.
        RuntimeError: If either object has not observed any data.
    """
    assert isinstance(other, Stats)
    assert self.prefix == other.prefix, "can't add two Stats objects with different prefixes"
    if self.stats is None or other.stats is None:
        raise RuntimeError("Cannot add Stats objects before both have been updated at least once.")

    new = Stats(prefix=self.prefix)
    new.n = int(self.n + other.n)

    w_self = self.n / new.n
    w_other = other.n / new.n
    new.stats = {
        "mean": float(self.stats["mean"]) * w_self + float(other.stats["mean"]) * w_other,
        "mean_squares": float(self.stats["mean_squares"]) * w_self
        + float(other.stats["mean_squares"]) * w_other,
        "max": float(max(float(self.stats["max"]), float(other.stats["max"]))),
        "min": float(min(float(self.stats["min"]), float(other.stats["min"]))),
        "mag": float(max(float(self.stats["mag"]), float(other.stats["mag"]))),
    }
    return new

def update(self, values: Union[np.ndarray, Iterable[float], float]):
    """Update the reservoir with a batch of values.

    Args:
        values: A scalar or array-like collection of numeric values. The
            input is flattened into a 1D array.
    """
    v = np.asarray(values, dtype=np.float32).ravel()
    m = int(v.size)
    if m == 0:
        return

    # Fill reservoir initially.
    if self.n < self.reservoir_size:
        take = min(self.reservoir_size - self.n, m)
        if take > 0:
            self.res = np.concatenate([self.res, v[:take]])
            self.n += int(take)
            v = v[take:]

    # Reservoir sampling for the remainder.
    for x in v:
        self.n += 1
        j = int(self.rng.integers(0, self.n))
        if j < self.reservoir_size:
            self.res[j] = x

def get(self) -> np.ndarray:
    """Return a copy of the reservoir buffer.

    Returns:
        A copy of the internal reservoir array (dtype ``np.float32``).
    """
    return self.res.copy()

def __add__(self, other: "Dist") -> "Dist":
    """Merge two reservoirs into a new reservoir.

    The merged reservoir is produced by sampling (without replacement) from
    the concatenation of both reservoirs, capped at :attr:`reservoir_size`.
    This is a pragmatic merge for logging/visualisation; it does not exactly
    reproduce a single-pass reservoir over the full underlying streams.

    Args:
        other: Another :class:`Dist` with the same :attr:`prefix` and
            :attr:`reservoir_size`.

    Returns:
        A new :class:`Dist` containing a merged reservoir sample.

    Raises:
        AssertionError: If :attr:`prefix` differs.
    """
    assert isinstance(other, Dist)
    assert self.prefix == other.prefix, "can't add two Dist objects with different prefixes"
    assert (
        self.reservoir_size == other.reservoir_size
    ), "can't add two Dist objects with different reservoir sizes"

    new = Dist(prefix=self.prefix, reservoir_size=self.reservoir_size, rng=0)
    new.n = int(self.n + other.n)

    both = np.concatenate([self.res, other.res])
    if both.size <= new.reservoir_size:
        new.res = both.astype(np.float32, copy=False)
    else:
        idx = np.random.default_rng(0).choice(
            both.size, size=new.reservoir_size, replace=False
        )
        new.res = both[idx].astype(np.float32, copy=False)
    return new