"""Dials that tie sampled values together.
Every linked value reads a dial between 0 and 1. A range maps the dial from its
first end to its second, so a reversed range falls as the dial rises. A
:class:`~oowl.sampling.Choice` splits the dial into one band per value, sized
by its probabilities in list order, and the linked values inside the chosen
value read the dial zoomed into that band. A dial can follow another one at a
rank correlation. Dials are turned once per scenario, or once per farm. A
resampled scenario turns its random dials afresh.
"""
from __future__ import annotations
import typing
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass
from typing import Any
import numpy as np
from ._random import as_rng, couple_quantiles, derive_seed
from ._repr import _describe
from ._sampling_specs import Choice, is_number, is_numeric_pair
__all__ = ["Link"]
_SCOPES = ("scenario", "farm")
_SELECTIONS = ("random", "balanced")
def _van_der_corput(index: int) -> float:
"""The ``index``-th point of the base-2 van der Corput sequence in [0, 1)."""
value, denominator = 0.0, 1.0
index = int(index)
while index:
denominator *= 2.0
index, remainder = divmod(index, 2)
value += remainder / denominator
return value
[docs]
class Link:
"""A dial that several sampled values read together.
``link(spec)`` returns a linked version of a ``(low, high)`` range or a
:class:`~oowl.sampling.Choice`, usable wherever a range or a choice is.
Values linked to one dial move together. A reversed range moves the other
way.
A scenario that ``on_error="resample"`` retries turns its random dials
again, as it redraws its other random values, while a balanced dial keeps
its place in the population. A farm retried inside a cluster keeps the
cluster's dials, which turn again when the whole cluster is retried.
Parameters
----------
follows
Another link this dial follows, or ``None``.
correlation
Rank correlation with ``follows``, in [-1, 1]. Requires ``follows``.
Defaults to 1, so the two dials turn together.
per
``"scenario"`` turns the dial once per scenario, so every farm of a
cluster reads the same value. ``"farm"`` turns it once per farm. A
value one cluster shares, its resource, surface, domain and farm or
site count, has no farm to read the dial for, so it cannot read a
per-farm link, and a per-scenario link cannot follow one.
selection
``"random"`` draws the dial independently per scenario. ``"balanced"``
spreads it evenly over scenarios, so linked choices meet their
probabilities and linked ranges are evenly covered.
"""
def __init__(
self,
*,
follows: Link | None = None,
correlation: float | None = None,
per: typing.Literal["scenario", "farm"] = "scenario",
selection: typing.Literal["random", "balanced"] = "random",
):
if follows is not None and not isinstance(follows, Link):
raise TypeError("follows must be a Link or None")
if correlation is not None and follows is None:
raise ValueError("correlation needs a link to follow")
if per not in _SCOPES:
raise ValueError("per must be 'scenario' or 'farm'")
if follows is not None and per == "scenario" and follows.per == "farm":
# One scenario dial cannot track a leader that differs per farm.
raise ValueError("a per-scenario link cannot follow a per-farm link")
correlation = 1.0 if correlation is None else float(correlation)
if not np.isfinite(correlation) or abs(correlation) > 1.0:
raise ValueError("correlation must be in [-1, 1]")
if selection not in _SELECTIONS:
raise ValueError("selection must be 'random' or 'balanced'")
self.follows = follows
self.correlation = correlation
self.per = per
self.selection = selection
def __call__(self, spec, *, scale: str | tuple | Callable = "linear") -> Linked:
"""Return ``spec`` read from this dial, a range scaled by ``scale``.
``scale`` is ``"linear"``, ``"log"`` for positive ranges spanning
orders of magnitude, ``("power", p)`` to bunch draws towards the first
end for ``p > 1``, or a function taking the dial in [0, 1] to [0, 1]
that only rises. A choice takes no scale.
"""
return Linked(self, spec, scale)
def __repr__(self) -> str:
# A leaderless link takes no correlation, so it prints none.
correlation = None if self.follows is None else self.correlation
return (
f"Link(follows={self.follows!r}, correlation={_describe(correlation)}, "
f"per={self.per!r}, selection={self.selection!r})"
)
def _validated_scale(scale, spec):
"""Check ``scale`` for ``spec`` and return it, a power as a float."""
if isinstance(spec, Choice):
if scale != "linear":
raise ValueError("a linked choice takes no scale")
return scale
if scale in ("linear", "log"):
if scale == "log" and not all(float(end) > 0 for end in spec):
raise ValueError("a log-scaled range needs two positive ends")
return scale
if isinstance(scale, tuple) and len(scale) == 2 and scale[0] == "power":
exponent = scale[1]
if not is_number(exponent):
raise TypeError(
f"a power scale exponent must be a number, not {type(exponent).__name__}"
)
if not np.isfinite(exponent) or exponent <= 0:
raise ValueError("a power scale needs a positive, finite exponent")
power = float(exponent)
return ("power", power)
if callable(scale):
return scale
raise ValueError("scale must be 'linear', 'log', ('power', p) or a function")
@dataclass(frozen=True, repr=False)
class Linked:
"""A range or choice read from a :class:`Link` dial."""
link: Link
spec: Any
scale: Any = "linear"
def __post_init__(self):
if not isinstance(self.link, Link):
raise TypeError("link must be a Link")
if isinstance(self.spec, Choice):
if self.spec.selection != "random":
raise ValueError(
"a linked choice follows its dial, so balance the dial "
"with Link(selection='balanced') instead"
)
elif not is_numeric_pair(self.spec):
raise TypeError("a link reads a (first, second) range or a Choice")
object.__setattr__(self, "scale", _validated_scale(self.scale, self.spec))
@property
def integer(self) -> bool:
"""Whether the spec is a range with two whole-number ends."""
return is_numeric_pair(self.spec) and all(
isinstance(end, (int, np.integer)) for end in self.spec
)
def ends(self) -> tuple[float, float]:
"""The range's two ends, in the order the dial reaches them."""
return float(self.spec[0]), float(self.spec[1])
def _position(self, dial: float) -> float:
"""The dial after scaling, in [0, 1]."""
if self.scale in ("linear", "log"):
return dial
if isinstance(self.scale, tuple):
return dial ** self.scale[1]
position = float(self.scale(dial))
if not np.isfinite(position) or not 0.0 <= position <= 1.0:
raise ValueError("a scale function must return a value in [0, 1]")
return position
def _between(self, first: float, second: float, dial: float) -> float:
position = self._position(dial)
if self.scale == "log":
return first * (second / first) ** position
return first + position * (second - first)
def value(self, dial: float) -> float | int:
"""The range's value at ``dial``. Whole-number ends give whole numbers."""
first, second = self.ends()
if not self.integer:
return float(self._between(first, second, dial))
# Each whole number owns an equal share of the dial.
if first <= second:
number = int(np.floor(self._between(first, second + 1.0, dial)))
return min(number, int(second))
number = int(np.ceil(self._between(first + 1.0, second, dial))) - 1
return max(number, int(second))
def band(self, dial: float) -> tuple[int, float, float]:
"""The chosen position and the ``(start, width)`` of its band."""
probabilities = np.asarray(self.spec.probabilities, dtype=float)
edges = np.concatenate([[0.0], np.cumsum(probabilities)])
position = int(np.searchsorted(edges, dial, side="right") - 1)
position = min(max(position, 0), len(probabilities) - 1)
# Skip empty bands so a zero-probability value is never picked.
while probabilities[position] == 0.0:
position -= 1
return position, float(edges[position]), float(probabilities[position])
def __repr__(self) -> str:
scale = "" if self.scale == "linear" else f", scale={self.scale!r}"
return f"{self.link!r}({self.spec!r}{scale})"
class _Dials:
"""The dials of one scenario, turned lazily and recorded."""
def __init__(self, numbers: Mapping[Link, int], index: int, seed: int, root: int):
self.numbers = numbers
self.index = int(index)
self.seed = int(seed)
self.root = int(root)
self.turned: dict[tuple[int, int | None], float] = {}
self.zooms: dict[int, list[tuple[float, float]]] = {}
def _number(self, link: Link) -> int:
try:
return self.numbers[link]
except KeyError:
raise ValueError(
"this Link is not part of the Sampler that is drawing it"
) from None
def raw(self, link: Link, farm: int) -> float:
"""The dial's value for ``farm``, before any zoom."""
number = self._number(link)
key = None if link.per == "scenario" else int(farm)
if (number, key) in self.turned:
return self.turned[number, key]
path = ("link", number) if key is None else ("link", number, "farm", key)
if link.selection == "balanced":
shift = as_rng(derive_seed(self.root, *path)).random()
dial = (_van_der_corput(self.index) + shift) % 1.0
else:
dial = float(as_rng(derive_seed(self.seed, *path)).random())
if link.follows is not None:
leading = self.raw(link.follows, farm)
dial = float(couple_quantiles(leading, dial, link.correlation))
self.turned[number, key] = dial
return dial
def read(self, link: Link, farm: int) -> float:
"""The dial as the value being drawn sees it, zoomed into any band."""
dial = self.raw(link, farm)
for start, width in self.zooms.get(id(link), ()):
dial = min(max((dial - start) / width, 0.0), 1.0)
return dial
@contextmanager
def zoom(self, link: Link, start: float, width: float) -> Iterator[None]:
"""Read ``link`` stretched over the band ``[start, start + width]``."""
stack = self.zooms.setdefault(id(link), [])
stack.append((start, width))
try:
yield
finally:
stack.pop()
def record(self) -> dict:
"""The turned dials by link number, and by farm for per-farm links."""
record: dict = {}
for (number, key), dial in sorted(
self.turned.items(), key=lambda item: (item[0][0], item[0][1] or 0)
):
if key is None:
record[str(number)] = dial
else:
record.setdefault(str(number), {})[str(key)] = dial
return record
_ACTIVE: ContextVar[_Dials | None] = ContextVar("oowl_dials", default=None)
@contextmanager
def scenario_dials(
numbers: Mapping[Link, int], index: int, seed: int, root: int
) -> Iterator[_Dials]:
"""Activate one scenario's dials, or keep the scenario already active."""
active = _ACTIVE.get()
if active is not None:
yield active
return
dials = _Dials(numbers, index, seed, root)
token = _ACTIVE.set(dials)
try:
yield dials
finally:
_ACTIVE.reset(token)
def _active() -> _Dials:
dials = _ACTIVE.get()
if dials is None:
raise ValueError("linked values are drawn by a Sampler")
return dials
def linked_value(spec: Linked, farm: int) -> float | int:
"""The value of a linked range for ``farm`` in the active scenario."""
return spec.value(_active().read(spec.link, farm))
def is_choice(spec) -> bool:
"""Whether ``spec`` is a :class:`~oowl.sampling.Choice`, linked or not."""
return isinstance(spec, Choice) or (
isinstance(spec, Linked) and isinstance(spec.spec, Choice)
)
@contextmanager
def chosen(
spec: Choice | Linked, index: int, seed: int | None
) -> Iterator[tuple[int, str, Any]]:
"""Select a value of ``spec`` for ``index`` as ``(position, selection, value)``.
A plain choice selects with ``seed`` unless it has its own. A linked choice
reads its dial instead and keeps the dial zoomed into the chosen band while
the value is drawn.
"""
if isinstance(spec, Choice):
position, value = spec.choose(index, seed=seed)
yield position, spec.selection, value
return
dials = _active()
position, start, width = spec.band(dials.read(spec.link, index))
with dials.zoom(spec.link, start, width):
yield position, "linked", spec.spec.values[position]
def _linked_items(spec, seen: set[int]) -> Iterator[Link | Linked]:
"""Yield every link and linked value inside ``spec``, outermost first."""
if id(spec) in seen:
return
seen.add(id(spec))
if isinstance(spec, Link):
yield spec
elif isinstance(spec, Linked):
yield spec
yield from _linked_items(spec.spec, seen)
elif isinstance(spec, Choice):
for item in spec.values:
yield from _linked_items(item, seen)
elif isinstance(spec, Mapping):
for item in spec.values():
yield from _linked_items(item, seen)
elif isinstance(spec, (tuple, list)):
for item in spec:
yield from _linked_items(item, seen)
def number_links(*specs) -> dict[Link, int]:
"""Number every link in ``specs`` in order of first appearance."""
# Keyed by the link itself, not ``id``, so a pickled sampler stays valid.
numbers: dict[Link, int] = {}
def visit_link(link: Link) -> None:
if link in numbers:
return
if link.follows is not None:
visit_link(link.follows)
numbers[link] = len(numbers)
seen: set[int] = set()
for spec in specs:
for item in _linked_items(spec, seen):
visit_link(item if isinstance(item, Link) else item.link)
return numbers
def per_farm_linked(spec) -> Linked | None:
"""The first value in ``spec`` read from a per-farm link, or ``None``."""
return next(
(
item
for item in _linked_items(spec, set())
if isinstance(item, Linked) and item.link.per == "farm"
),
None,
)
