kalfa 3.0.0.dev0__py3-none-any.whl

kalfa/commands/sweep/strategies.py

@@ -0,0 +1,228 @@
+import copy
+import difflib
+import itertools
+import json
+
+import numpy
+
+from kalfa.commands.sweep.space import draw, is_int, space_values, spread
+from kalfa.core import aliases
+from kalfa.core.errors import BuildError
+from kalfa.records import reader
+from kalfa.records.sweep import mode_of
+from kalfa.std.common import positive_int
+
+
+SAMPLERS = ("tpe", "random")
+OPTUNA_MISSING = "optuna needs the optuna package (uv sync installs the dependencies of kalfa)"
+STRATEGY_HINT = "grid, random, optuna or {use: <name or import path>, <param>: ...}"
+
+
+def seed_value(seed):
+    return None if seed is None else int(seed)
+
+
+def resumed_seed(seed, count):
+    return None if seed is None else int(numpy.random.SeedSequence([seed, count]).generate_state(1)[0])
+
+
+class Grid:
+    def __init__(self, steps=None):
+        self.steps = None if steps is None else positive_int(steps, "grid steps")
+
+    def axes(self, space):
+        axes = {}
+        for name, (kind, values) in space_values(space).items():
+            if kind == "list":
+                axes[name] = values
+            elif self.steps is None:
+                raise BuildError(f"sweep.space.{name}: the {kind} range needs steps on the grid, as in "
+                                 f"strategy: {{use: grid, steps: 5}}, or a list of values")
+            else:
+                axes[name] = spread(kind, values, self.steps)
+        return axes
+
+    def points(self, space, seed=0):
+        axes = self.axes(space)
+        names = list(axes)
+        return [dict(zip(names, copy.deepcopy(combination))) for combination in itertools.product(*axes.values())]
+
+
+class Random:
+    def __init__(self, trials, seed=None):
+        self.trials = positive_int(trials, "random trials")
+        self.seed = seed_value(seed)
+
+    def points(self, space, seed=0):
+        dimensions = space_values(space)
+        generator = numpy.random.default_rng(self.seed if self.seed is not None else seed)
+        found, seen = [], set()
+        for _ in range(self.trials * 10):
+            point = {name: draw(generator, kind, values) for name, (kind, values) in dimensions.items()}
+            key = json.dumps(point, sort_keys=True, default=str)
+            if key in seen:
+                continue
+            seen.add(key)
+            found.append(point)
+            if len(found) == self.trials:
+                break
+        return found
+
+
+def optuna_module():
+    try:
+        import optuna
+    except ImportError as exc:
+        raise BuildError(OPTUNA_MISSING) from exc
+    return optuna
+
+
+class Optuna:
+    def __init__(self, trials, sampler="tpe", seed=None):
+        self.trials = positive_int(trials, "optuna trials")
+        if sampler not in SAMPLERS:
+            raise BuildError(f"the optuna sampler is tpe or random, got {sampler!r}")
+        self.sampler = sampler
+        self.seed = seed_value(seed)
+        self.study = None
+        optuna_module()
+
+    def start(self, objective, seed=None):
+        optuna = optuna_module()
+        optuna.logging.set_verbosity(optuna.logging.WARNING)
+        maker = optuna.samplers.TPESampler if self.sampler == "tpe" else optuna.samplers.RandomSampler
+        direction = "maximize" if mode_of(objective) == "max" else "minimize"
+        self.study = optuna.create_study(direction=direction, sampler=maker(seed=self.seed if seed is None else seed))
+        return self.study
+
+    def resume(self, space, objective, points):
+        optuna = optuna_module()
+        points = list(points)
+        study = self.start(objective, resumed_seed(self.seed, len(points)))
+        dimensions = space_values(space)
+        told = 0
+        for values, value in points:
+            trial = finished_trial(optuna, dimensions, values, value)
+            if trial is not None:
+                study.add_trial(trial)
+                told += 1
+        return told
+
+    def ask(self, space, objective):
+        dimensions = space_values(space)
+        study = self.study if self.study is not None else self.start(objective)
+        trial = study.ask()
+        return trial, {name: suggest(trial, name, kind, values) for name, (kind, values) in dimensions.items()}
+
+    def tell(self, trial, value):
+        optuna = optuna_module()
+        if value is None or not reader.is_finite(value):
+            self.study.tell(trial, state=optuna.trial.TrialState.FAIL)
+        else:
+            self.study.tell(trial, float(value))
+
+
+def is_plain(value):
+    return value is None or isinstance(value, (bool, int, float, str))
+
+
+def suggest(trial, name, kind, values):
+    if kind == "list":
+        if all(is_plain(value) for value in values):
+            return trial.suggest_categorical(name, values)
+        return copy.deepcopy(values[trial.suggest_categorical(name, list(range(len(values))))])
+    low, high = values
+    if kind == "int":
+        return int(trial.suggest_int(name, low, high))
+    return float(trial.suggest_float(name, low, high, log=kind == "log"))
+
+
+def distribution_of(optuna, kind, values, value):
+    found = optuna.distributions
+    if kind == "list":
+        if all(is_plain(choice) for choice in values):
+            return (found.CategoricalDistribution(values), value) if value in values else None
+        index = next((index for index, choice in enumerate(values) if choice == value), None)
+        return None if index is None else (found.CategoricalDistribution(list(range(len(values)))), index)
+    low, high = values
+    if not reader.is_finite(value) or not low <= value <= high:
+        return None
+    if kind == "int":
+        return (found.IntDistribution(low, high), int(value)) if is_int(value) else None
+    return found.FloatDistribution(low, high, log=kind == "log"), float(value)
+
+
+def finished_trial(optuna, dimensions, values, value):
+    if not isinstance(values, dict):
+        return None
+    params, distributions = {}, {}
+    for name, (kind, choices) in dimensions.items():
+        found = distribution_of(optuna, kind, choices, values[name]) if name in values else None
+        if found is None:
+            return None
+        distributions[name], params[name] = found
+    try:
+        if value is None or not reader.is_finite(value):
+            return optuna.trial.create_trial(params=params, distributions=distributions,
+                                             state=optuna.trial.TrialState.FAIL)
+        return optuna.trial.create_trial(params=params, distributions=distributions, value=float(value))
+    except ValueError:
+        return None
+
+
+STRATEGIES = {"grid": Grid, "random": Random, "optuna": Optuna}
+
+
+def strategy_call(raw):
+    if raw is None:
+        return "grid", {}
+    if isinstance(raw, str) and raw:
+        return raw, {}
+    if isinstance(raw, dict) and isinstance(raw.get("use"), str) and raw["use"]:
+        return raw["use"], {str(key): value for key, value in raw.items() if key != "use"}
+    raise BuildError(f"sweep.strategy is {STRATEGY_HINT}, got {raw!r}")
+
+
+def strategy_names():
+    return [*STRATEGIES, *(entry.name for entry in aliases.catalog(slot="strategy") if entry.name not in STRATEGIES)]
+
+
+def strategy_target(use):
+    if use in STRATEGIES:
+        return STRATEGIES[use]
+    entry = aliases.entry_of(use)
+    if entry is not None:
+        if entry.slot not in (None, "strategy"):
+            raise BuildError(f"{use!r} is a {entry.slot}, not a sweep strategy; the strategies are "
+                             + ", ".join(strategy_names()))
+        return aliases.import_target(entry.target)
+    if aliases.is_import_path(use):
+        return aliases.import_target(use)
+    close = difflib.get_close_matches(use, strategy_names(), n=3, cutoff=0.6)
+    hint = f"; did you mean {', '.join(close)}" if close else ""
+    raise BuildError(f"{use!r} is not a sweep strategy; the strategies are {', '.join(strategy_names())} or an "
+                     f"import path such as pkg.module.Class{hint}")
+
+
+def strategy_of(raw):
+    use, options = strategy_call(raw)
+    target = strategy_target(use)
+    try:
+        strategy = target(**options)
+    except BuildError:
+        raise
+    except Exception as exc:
+        raise BuildError(f"sweep.strategy {use}: {type(exc).__name__}: {exc}") from exc
+    if callable(getattr(strategy, "points", None)):
+        return strategy
+    if not (callable(getattr(strategy, "ask", None)) and callable(getattr(strategy, "tell", None))):
+        raise BuildError(f"sweep.strategy {use}: a strategy has points(space, seed), or ask(space, objective) and "
+                         f"tell(trial, value)")
+    positive_int(getattr(strategy, "trials", None), f"sweep.strategy {use}: trials")
+    return strategy
+
+
+def strategy_note(raw):
+    if isinstance(raw, dict):
+        return copy.deepcopy(raw)
+    return {"use": raw if isinstance(raw, str) else "grid"}

kalfa/core/aliases.py

@@ -0,0 +1,229 @@
+import difflib
+import importlib
+import importlib.util
+import inspect
+import logging
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from .errors import BuildError, error
+
+
+LOGGER = logging.getLogger("kalfa.aliases")
+HINT = "kalfa legos lists the names"
+
+
+@dataclass
+class Entry:
+    name: str
+    target: str
+    slot: str | None = None
+    doc: str = ""
+    params: dict = field(default_factory=dict)
+    examples: list = field(default_factory=list)
+    plugin: bool = False
+
+
+ALIASES = {}
+SEARCH_DIRS = []
+
+
+def target_path(target):
+    if isinstance(target, str):
+        return target
+    if inspect.ismodule(target):
+        return target.__name__
+    module = getattr(target, "__module__", None)
+    qualname = getattr(target, "__qualname__", None) or getattr(target, "__name__", None)
+    if module and qualname:
+        return f"{module}.{qualname}"
+    return repr(target)
+
+
+def imports_as(path, target):
+    try:
+        return import_target(path) is target
+    except BuildError:
+        return False
+
+
+def same_target(first, second):
+    if first is second:
+        return True
+    if isinstance(first, str) != isinstance(second, str):
+        path, target = (first, second) if isinstance(first, str) else (second, first)
+        if imports_as(path, target):
+            return True
+    if not isinstance(first, str) and not isinstance(second, str):
+        if "<" in target_path(first) or "<" in target_path(second):
+            return False
+    return target_path(first) == target_path(second)
+
+
+def register(name, target, *, slot=None, doc="", params=None, examples=None, plugin=False):
+    if not isinstance(name, str) or not name:
+        raise BuildError(f"an alias name must be a non empty string, got {name!r}")
+    if target is None:
+        raise BuildError(f"{name!r}: the target must be an import path or an object")
+    entry = Entry(name, target, slot, doc or "", dict(params or {}), list(examples or []), bool(plugin))
+    existing = ALIASES.get(name)
+    if existing is not None:
+        if same_target(existing.target, target):
+            if doc and not existing.doc:
+                existing.doc = doc
+            if slot and not existing.slot:
+                existing.slot = slot
+            return existing
+        if existing.plugin and plugin:
+            LOGGER.warning("alias %r moves from %s to %s", name, target_path(existing.target), target_path(target))
+        else:
+            raise BuildError(f"{name!r} is already registered as {target_path(existing.target)}; "
+                             f"{target_path(target)} cannot take the name ({HINT})")
+    ALIASES[name] = entry
+    return entry
+
+
+def annotate(name, params=None, examples=None):
+    entry = ALIASES.get(name) if isinstance(name, str) else None
+    if entry is None:
+        raise BuildError(f"{name!r} is not a registered name, so it takes no notes ({HINT})")
+    entry.params.update(params or {})
+    for example in examples or ():
+        if example not in entry.examples:
+            entry.examples.append(example)
+    return entry
+
+
+def lego(name, *, slot=None, doc="", params=None, examples=None):
+    def decorate(target):
+        text = doc or (getattr(target, "__doc__", None) or "").strip()
+        register(name, target, slot=slot, doc=text, params=params, examples=examples, plugin=True)
+        return target
+    return decorate
+
+
+def is_import_path(text):
+    return isinstance(text, str) and "." in text and not text.startswith(".")
+
+
+def resolve(name):
+    entry = ALIASES.get(name) if isinstance(name, str) else None
+    if entry is not None:
+        return entry.target
+    if is_import_path(name):
+        return name
+    raise BuildError(f"{name!r} is not a registered name and not an import path such as pkg.module.Class ({HINT})")
+
+
+def suggestion(name):
+    close = difflib.get_close_matches(str(name), list(ALIASES), n=3, cutoff=0.6)
+    return "did you mean " + ", ".join(close) if close else None
+
+
+def add_search_dir(directory):
+    if directory and str(directory) not in SEARCH_DIRS:
+        SEARCH_DIRS.append(str(directory))
+        if str(directory) not in sys.path:
+            sys.path.insert(0, str(directory))
+
+
+def missing_prefix(exc, module_name):
+    missing = getattr(exc, "name", None)
+    return bool(missing) and (missing == module_name or module_name.startswith(missing + "."))
+
+
+def import_target(path):
+    if not isinstance(path, str):
+        return path
+    if not path or path.startswith(".") or path.endswith("."):
+        raise BuildError(f"{path!r} is not an import path such as pkg.module.Class")
+    parts = path.split(".")
+    if parts[0] not in sys.modules:
+        importlib.invalidate_caches()
+    failure = None
+    for cut in range(len(parts), 0, -1):
+        module_name = ".".join(parts[:cut])
+        try:
+            module = importlib.import_module(module_name)
+        except ModuleNotFoundError as exc:
+            if not missing_prefix(exc, module_name):
+                raise BuildError(f"{path}: importing {module_name} failed: {exc}") from exc
+            failure = exc
+            continue
+        except Exception as exc:
+            raise BuildError(f"{path}: importing {module_name} failed: {type(exc).__name__}: {exc}") from exc
+        target = module
+        for attribute in parts[cut:]:
+            try:
+                target = getattr(target, attribute)
+            except AttributeError as exc:
+                raise BuildError(f"{path}: {target_path(target)} has no attribute {attribute!r}") from exc
+        return target
+    raise BuildError(f"{path}: {failure}")
+
+
+def plugin_file(name):
+    candidate = Path(name)
+    if candidate.suffix != ".py":
+        return None
+    roots = [Path.cwd(), *[Path(directory) for directory in SEARCH_DIRS]]
+    if candidate.is_absolute():
+        roots = [None]
+    for root in roots:
+        path = candidate if root is None else root / candidate
+        if path.is_file():
+            return path.resolve()
+    return None
+
+
+def import_file(path):
+    name = path.stem
+    found = sys.modules.get(name)
+    if found is not None and getattr(found, "__file__", None) and Path(found.__file__).resolve() == path:
+        return found
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    try:
+        spec.loader.exec_module(module)
+    except BaseException:
+        sys.modules.pop(name, None)
+        raise
+    return module
+
+
+def import_plugins(names, problems):
+    loaded = []
+    importlib.invalidate_caches()
+    for name in names or ():
+        try:
+            path = plugin_file(str(name))
+            if path is not None:
+                loaded.append(import_file(path))
+            elif isinstance(name, str) and name.endswith(".py"):
+                raise ImportError(f"no file {name} under the working directory or the config directories")
+            else:
+                loaded.append(importlib.import_module(str(name)))
+        except Exception as exc:
+            problems.append(error("plugin_import", f"{name}: {type(exc).__name__}: {exc}", "plugins",
+                                  "a plugin is a module on sys.path or a .py file next to the config"))
+    return loaded
+
+
+def catalog(slot=None, word=None):
+    needle = word.lower() if isinstance(word, str) and word else None
+    found = []
+    for name in sorted(ALIASES):
+        entry = ALIASES[name]
+        if slot is not None and entry.slot != slot:
+            continue
+        if needle is not None and needle not in name.lower() and needle not in entry.doc.lower():
+            continue
+        found.append(entry)
+    return found
+
+
+def entry_of(name):
+    return ALIASES.get(name)
+