"""Specification objects used by population sampling."""
from __future__ import annotations
import threading
import typing
from array import array
from collections import OrderedDict
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any
import numpy as np
from ._random import as_rng, derive_seed, pin_seed, resolve_seed
from ._repr import _ArgumentRepr
__all__ = ["Choice", "Literal"]
def _probabilities(values, count: int) -> tuple[float, ...]:
"""Validate one probability per choice, uniform when ``values`` is ``None``."""
if values is None:
result = np.full(count, 1.0 / count)
else:
result = np.asarray(values, dtype=float)
if result.shape != (count,):
raise ValueError("probabilities must have one value per choice")
if not np.all(np.isfinite(result)) or np.any(result < 0):
raise ValueError("probabilities must be finite and non-negative")
# Match NumPy's chooser tolerance.
if abs(float(result.sum()) - 1.0) > np.sqrt(np.finfo(np.float64).eps):
raise ValueError("probabilities must sum to 1")
return tuple(float(value) for value in result)
#: Balanced schedules kept, so a population draws each prefix only once.
_SCHEDULE_CACHE_SIZE = 64
_schedules: OrderedDict[tuple, tuple[array, list[int]]] = OrderedDict()
_schedules_lock = threading.Lock()
def _balanced_choice(
probabilities: tuple[float, ...], root_seed: int, index: int
) -> int:
"""The position a balanced schedule rooted at ``root_seed`` picks at ``index``.
Each step picks the value furthest below its share, breaking exact ties
with the step's own seed. A schedule depends only on its probabilities
and root, so its prefix is cached and extended rather than replayed.
"""
key = (probabilities, root_seed)
with _schedules_lock:
schedule = _schedules.pop(key, None) or (array("I"), [0] * len(probabilities))
_schedules[key] = schedule
if len(_schedules) > _SCHEDULE_CACHE_SIZE:
_schedules.popitem(last=False)
choices, counts = schedule
for step in range(len(choices), index + 1):
deficits = [
(step + 1) * share - count
for share, count in zip(probabilities, counts)
]
top = max(deficits)
best = [
position
for position, deficit in enumerate(deficits)
if abs(deficit - top) <= 1e-12
]
choice = best[0]
if len(best) > 1:
tie_seed = derive_seed(root_seed, step)
choice = int(as_rng(tie_seed).choice(np.asarray(best)))
counts[choice] += 1
choices.append(choice)
return choices[index]
[docs]
@dataclass(frozen=True, repr=False)
class Choice(_ArgumentRepr):
"""A weighted choice among values, made separately for each index drawn.
``values`` may hold plain objects, constructors, recipes
``(constructor, {kwargs})``, or mappings of farm configuration overrides.
``probabilities`` defaults to uniform. ``selection="random"`` draws each
index independently, and ``"balanced"`` follows a schedule that keeps every
prefix of a population close to the probabilities. ``seed`` fixes the
selection whatever sampler draws it, and ``None`` follows the sampler.
"""
values: tuple[Any, ...]
probabilities: tuple[float, ...]
selection: typing.Literal["random", "balanced"] = "random"
seed: int | np.random.Generator | None = None
def __init__(
self,
values: Sequence[Any],
probabilities: Sequence[float] | None = None,
*,
selection: typing.Literal["random", "balanced"] = "random",
seed: int | np.random.Generator | None = None,
):
if isinstance(values, (str, bytes)):
raise TypeError("values must be a non-string sequence")
try:
values = tuple(values)
except TypeError as exc:
raise TypeError("values must be a non-empty sequence") from exc
if not values:
raise ValueError("values must contain at least one choice")
if selection not in ("random", "balanced"):
raise ValueError("selection must be 'random' or 'balanced'")
object.__setattr__(self, "values", values)
object.__setattr__(
self, "probabilities", _probabilities(probabilities, len(values))
)
object.__setattr__(self, "selection", selection)
object.__setattr__(self, "seed", pin_seed(seed))
[docs]
def choose(self, index: int = 0, *, seed: int | None = None) -> tuple[int, Any]:
"""Return ``(position, value)`` of the member chosen for ``index``.
``seed`` roots the selection when the choice has no seed of its own.
With neither, the selection draws fresh entropy.
"""
if not is_number(index):
raise TypeError(
f"choice index must be an integer, not {type(index).__name__}"
)
if not is_whole(index) or index < 0:
raise ValueError("choice index must be a non-negative integer")
root_seed = resolve_seed(self.seed if self.seed is not None else seed)
if self.selection == "random":
selection_seed = derive_seed(root_seed, "choice", int(index))
choice = int(
as_rng(selection_seed).choice(len(self.values), p=self.probabilities)
)
else:
choice = _balanced_choice(self.probabilities, root_seed, int(index))
return choice, self.values[choice]
[docs]
@dataclass(frozen=True)
class Literal:
"""A value the sampler passes through unchanged rather than drawing from."""
value: Any
def is_recipe(value: object) -> bool:
"""Whether *value* has the structural form ``(constructor, mapping)``."""
if not isinstance(value, tuple) or len(value) != 2:
return False
constructor, kwargs = value
return isinstance(kwargs, Mapping) and (
isinstance(constructor, Choice) or callable(constructor)
)
def is_distribution(value: object) -> bool:
"""Whether *value* is a frozen ``scipy.stats`` distribution to draw from."""
return hasattr(value, "rvs") and hasattr(value, "dist")
def is_number(value: object) -> bool:
"""Whether *value* is a real number, booleans excluded."""
return isinstance(value, (int, float, np.integer, np.floating)) and not isinstance(
value, (bool, np.bool_)
)
def is_whole(value: object) -> bool:
"""Whether *value* is a whole number, booleans excluded."""
return isinstance(value, (int, np.integer)) and not isinstance(
value, (bool, np.bool_)
)
def is_numeric_pair(value: object) -> bool:
"""Whether *value* is a tuple of two numbers, a range candidate."""
return (
isinstance(value, tuple)
and len(value) == 2
and all(is_number(item) for item in value)
)
def as_real(value, name: str) -> float:
"""Convert a drawn ``value`` to ``float``, or raise a ``TypeError`` naming it."""
if not isinstance(value, (str, bytes, bool, np.bool_)):
try:
return float(value)
except (TypeError, ValueError):
pass
raise TypeError(f"{name} must draw a real number, got {type(value).__name__}")
def nested_choice_seed(choice_seed: int | None) -> int | None:
"""The seed a choice nested in a chosen value selects with.
Deriving it keeps the inner selection independent of the outer one.
"""
return None if choice_seed is None else derive_seed(choice_seed, "selected")
def attempt_choice_seed(spec, choice_seed: int | None, attempt: int) -> int | None:
"""The seed scalar ``spec`` selects with on ``attempt`` of a retried draw.
A retry reselects a random choice, as it redraws a component, so a value
that cannot be drawn is not retried until the attempts run out. A balanced
choice keeps its place in the population schedule, and a pinned or linked
choice ignores the seed. The first attempt selects with ``choice_seed``.
"""
if (
attempt == 0
or choice_seed is None
or not isinstance(spec, Choice)
or spec.selection != "random"
):
return choice_seed
return derive_seed(choice_seed, "attempt", int(attempt))
def draw_scalar(
spec,
rng: np.random.Generator,
name: str,
*,
index: int = 0,
choice_seed: int | None = None,
minimum: float | None = None,
attempt: int = 0,
) -> float:
"""Resolve one scalar from a specification.
A specification is a number, a ``(low, high)`` uniform range, a frozen
``scipy.stats`` distribution such as ``loguniform(0.03, 1.0)``, a callable
``f(rng) -> number``, a :class:`~oowl.sampling.Literal` number, a
:class:`~oowl.sampling.Choice` of those, or a range or choice read from a
:class:`~oowl.sampling.Link`. ``minimum`` is an inclusive lower bound.
``None`` requires a value above zero. ``attempt`` counts earlier failed
tries of the draw, see :func:`attempt_choice_seed`.
"""
from ._links import Linked, chosen, is_choice, linked_value
if is_choice(spec):
with chosen(spec, index, attempt_choice_seed(spec, choice_seed, attempt)) as (
_,
_,
selected,
):
return draw_scalar(
selected,
rng,
name,
index=index,
choice_seed=nested_choice_seed(choice_seed),
minimum=minimum,
attempt=attempt,
)
if isinstance(spec, Literal):
value = spec.value
elif isinstance(spec, Linked):
value = linked_value(spec, index)
elif is_distribution(spec):
value = spec.rvs(random_state=rng)
elif callable(spec) and not isinstance(spec, type):
value = spec(rng)
elif is_number(spec):
value = spec
elif is_numeric_pair(spec):
low, high = (float(item) for item in spec)
if low > high:
raise ValueError(f"{name} range must have low <= high")
value = rng.uniform(low, high)
else:
raise TypeError(
f"{name} must be a number, a numeric range, a scipy.stats "
"distribution, or f(rng) -> number"
)
value = as_real(value, name)
if not np.isfinite(value):
raise ValueError(f"{name} must be finite")
if minimum is None:
if value <= 0:
raise ValueError(f"{name} must be finite and > 0")
elif value < minimum:
raise ValueError(f"{name} must be finite and >= {minimum:g}")
return value
def describe_scalar(spec) -> Any:
"""Return a JSON-compatible description of a scalar specification."""
from ._links import Linked
if isinstance(spec, Linked):
return {"linked": describe_scalar(spec.spec)}
if isinstance(spec, Literal):
return describe_scalar(spec.value)
if isinstance(spec, Choice):
return {"choice": [describe_scalar(value) for value in spec.values]}
if is_numeric_pair(spec):
return [float(item) for item in spec]
if is_distribution(spec):
return {
"distribution": spec.dist.name,
"args": [float(value) for value in spec.args],
"kwds": {key: float(value) for key, value in spec.kwds.items()},
}
if is_number(spec):
return float(spec)
if callable(spec):
return getattr(spec, "__name__", type(spec).__name__)
return type(spec).__name__
def _json_safe(value):
"""Return a JSON-compatible copy of a metadata value."""
if isinstance(value, dict):
return {str(key): _json_safe(item) for key, item in value.items()}
if isinstance(value, (list, tuple)):
return [_json_safe(item) for item in value]
if isinstance(value, np.ndarray):
return _json_safe(value.tolist())
if isinstance(value, np.bool_):
return bool(value)
if isinstance(value, np.integer):
return int(value)
if isinstance(value, np.floating):
return float(value)
if isinstance(value, (str, int, float, bool)) or value is None:
return value
return str(value)
