Source code for psp._control

"""Validation of the ``control`` settings of :func:`psp.PSP`."""

from __future__ import annotations

import math
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any

import torch

__all__ = ["Control", "UNBOUNDED"]

#: Value of ``iterations`` or ``population`` that lifts that stopping rule.
#: ``0``, ``None`` and ``inf`` are all accepted as "unbounded", as in R where ``0`` is used.
UNBOUNDED = 0

_INT_MAX = 2**63 - 1


def _is_unbounded(value: Any) -> bool:
    if value is None:
        return True
    if isinstance(value, float) and math.isinf(value):
        return True
    return int(value) == 0


def _as_int(name: str, value: Any) -> int:
    if _is_unbounded(value):
        return _INT_MAX
    if isinstance(value, bool) or not float(value).is_integer():
        raise ValueError(f"control['{name}'] must be a positive integer, 0 or None, got {value!r}")
    value = int(value)
    if value < 0:
        raise ValueError(f"control['{name}'] must be a positive integer, 0 or None, got {value!r}")
    return value


def _as_vector(name: str, value: Any, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    tensor = torch.as_tensor(value, dtype=dtype, device=device).reshape(-1)
    if tensor.numel() == 0:
        raise ValueError(f"control['{name}'] must not be empty")
    return tensor


def _as_names(name: str, value: Any, prefix: str, count: int) -> tuple[str, ...]:
    if value is None:
        return tuple(f"{prefix}_{i + 1}" for i in range(count))
    if isinstance(value, str):
        value = [value]
    names = tuple(str(v) for v in value)
    return names


[docs] @dataclass(frozen=True) class Control: """Validate control arguments for the parameter space partitioning run. It checks that the required keys are present and that the values are of the right type and shape. Use ``Control.from_mapping({"lower": [0, 0], "upper": [1, 1], ...})``, or pass the mapping to :func:`psp.PSP` directly. Attributes ---------- init Starting parameter sets, one per row, shape ``(starts, dimensions)``. lower, upper Bounds of each parameter, shape ``(dimensions,)``. Use ``-inf``/``inf`` for an unbounded parameter. radius Largest step the search takes in any one parameter, shape ``(dimensions,)``. population Number of parameter sets each pattern needs before the search stops proposing from it. ``2**63 - 1`` when unbounded. iterations Cap on the number of iterations. ``2**63 - 1`` when unbounded. parameter_names, stimuli_names Column names of the files written with ``save=True``. dimensionality Size of the square matrix ``discretize`` returns, or ``None`` to take it from the first evaluation. responses Length of the vector ``model`` returns, or ``None`` to take it from the first evaluation. """ init: torch.Tensor lower: torch.Tensor upper: torch.Tensor radius: torch.Tensor population: int iterations: int parameter_names: tuple[str, ...] stimuli_names: tuple[str, ...] | None dimensionality: int | None responses: int | None @property def dimensions(self) -> int: """Number of free parameters.""" return int(self.init.shape[1]) @property def population_unbounded(self) -> bool: return self.population == _INT_MAX @property def iterations_unbounded(self) -> bool: return self.iterations == _INT_MAX
[docs] @classmethod def from_mapping( cls, control: Mapping[str, Any] | Control, *, device: torch.device | str | None = None, dtype: torch.dtype | None = None, ) -> Control: """Validate a ``control`` mapping with the expected keys. Required keys: ``init``, ``lower``, ``upper``, ``radius``, and at least one of ``population`` and ``iterations``. Optional keys: ``parameter_names``, ``stimuli_names``, ``dimensionality``, ``responses``. Raises ------ ValueError If any of the required keys are missing, or if any of the values are of the wrong type or shape. """ if isinstance(control, Control): if device is None and dtype is None: return control control = control.to_mapping() if not isinstance(control, Mapping): raise TypeError( f"control must be a mapping with the expected keys, not {type(control).__name__}" ) unknown = set(control) - _KNOWN_KEYS if unknown: raise ValueError(f"unknown control keys: {sorted(unknown)}") for key in ("init", "lower", "upper", "radius"): if key not in control: raise ValueError(f"control['{key}'] is required") iterations = _as_int("iterations", control.get("iterations", UNBOUNDED)) population = _as_int("population", control.get("population", UNBOUNDED)) if iterations == _INT_MAX and population == _INT_MAX: raise ValueError( "A reasonable threshold must be set by either adjusting iterations or population." ) init_value = control["init"] if dtype is None: dtype = init_value.dtype if torch.is_tensor(init_value) else torch.float64 if not dtype.is_floating_point: dtype = torch.float64 if device is None: device = init_value.device if torch.is_tensor(init_value) else torch.device("cpu") device = torch.device(device) init = torch.as_tensor(init_value, dtype=dtype, device=device) if init.ndim == 1: init = init.unsqueeze(0) if init.ndim != 2 or init.numel() == 0: raise ValueError("control['init'] must be a matrix with one parameter set per row") dimensions = int(init.shape[1]) lower = _as_vector("lower", control["lower"], device, dtype) upper = _as_vector("upper", control["upper"], device, dtype) if lower.numel() != dimensions or upper.numel() != dimensions: raise ValueError("init, lower and upper must have the same length.") if bool((lower > upper).any()): raise ValueError("control['lower'] must not exceed control['upper']") radius = _as_vector("radius", control["radius"], device, dtype) if radius.numel() == 1: radius = radius.expand(dimensions).clone() elif radius.numel() != dimensions: raise ValueError( "control['radius'] must be a single number or one number per parameter" ) if bool((radius <= 0).any()) or not bool(torch.isfinite(radius).all()): raise ValueError("control['radius'] must be positive and finite") parameter_names = _as_names( "parameter_names", control.get("parameter_names"), "parameter", dimensions ) if len(parameter_names) != dimensions: raise ValueError("Length of param_names must equal to the number of dimensions") responses = control.get("responses") if responses is not None: responses = _as_int("responses", responses) if responses in (0, _INT_MAX): raise ValueError("control['responses'] must be a positive integer") dimensionality = control.get("dimensionality") if dimensionality is not None: dimensionality = _as_int("dimensionality", dimensionality) if dimensionality in (0, _INT_MAX): raise ValueError("control['dimensionality'] must be a positive integer") stimuli = control.get("stimuli_names") stimuli_names: tuple[str, ...] | None if stimuli is None: stimuli_names = ( None if responses is None else _as_names("stimuli_names", None, "stimulus", responses) ) else: stimuli_names = _as_names("stimuli_names", stimuli, "stimulus", 0) if responses is not None and len(stimuli_names) != responses: raise ValueError("Length of stimuli_names must equal to control['responses']") return cls( init=init, lower=lower, upper=upper, radius=radius, population=population, iterations=iterations, parameter_names=parameter_names, stimuli_names=stimuli_names, dimensionality=dimensionality, responses=responses, )
[docs] def to_mapping(self) -> dict[str, Any]: """The settings as a plain dict.""" return { "init": self.init, "lower": self.lower, "upper": self.upper, "radius": self.radius, "population": UNBOUNDED if self.population_unbounded else self.population, "iterations": UNBOUNDED if self.iterations_unbounded else self.iterations, "parameter_names": list(self.parameter_names), "stimuli_names": None if self.stimuli_names is None else list(self.stimuli_names), "dimensionality": self.dimensionality, "responses": self.responses, }
_KNOWN_KEYS = { "init", "lower", "upper", "radius", "population", "iterations", "parameter_names", "stimuli_names", "dimensionality", "responses", }