Jarvis-HEP-Sampler 1.0.0__py3-none-any.whl

jarvis_hep_sampler/__init__.py

@@ -0,0 +1,141 @@
+"""Jarvis-HEP-Sampler: the plug-in SDK for user-defined Jarvis-HEP samplers.
+
+A plug-in is one ``.py`` file with a ``PLUGIN`` dict, an optional
+``OPTIONS_SCHEMA`` and a few functions (``prepare`` / ``draw`` / hooks).
+Users create it with ``Jarvis sampler init``, load it once with
+``Jarvis sampler load``, and then write ``Sampling.Method: <Name>`` like any
+built-in method.
+
+This package owns the plug-in format: metadata rules, loading, the user
+registry, random streams, output checks, the logging facade and the offline
+test kit. Jarvis-HEP owns the runtime. Names are imported lazily so reading
+the registry stays light.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import TYPE_CHECKING, Any
+
+from jarvis_hep_sampler._version import __version__
+
+_LAZY_ATTRS: dict[str, tuple[str, str]] = {
+    # metadata
+    "API_VERSION": ("jarvis_hep_sampler.metadata", "API_VERSION"),
+    "PluginMetadata": ("jarvis_hep_sampler.metadata", "PluginMetadata"),
+    "check_method_name": ("jarvis_hep_sampler.metadata", "check_method_name"),
+    # contexts and logging
+    "BlockContext": ("jarvis_hep_sampler.contexts", "BlockContext"),
+    "DrawContext": ("jarvis_hep_sampler.contexts", "DrawContext"),
+    "SamplerStats": ("jarvis_hep_sampler.contexts", "SamplerStats"),
+    "SetupContext": ("jarvis_hep_sampler.contexts", "SetupContext"),
+    "LogBackend": ("jarvis_hep_sampler.log", "LogBackend"),
+    "SamplerLog": ("jarvis_hep_sampler.log", "SamplerLog"),
+    # errors
+    "PluginConfigError": ("jarvis_hep_sampler.errors", "PluginConfigError"),
+    "PluginError": ("jarvis_hep_sampler.errors", "PluginError"),
+    "PluginLoadError": ("jarvis_hep_sampler.errors", "PluginLoadError"),
+    "PluginOutputError": ("jarvis_hep_sampler.errors", "PluginOutputError"),
+    "PluginResumeMismatch": ("jarvis_hep_sampler.errors", "PluginResumeMismatch"),
+    "RegistryError": ("jarvis_hep_sampler.errors", "RegistryError"),
+    # streams
+    "STREAM_PROTOCOL": ("jarvis_hep_sampler.streams", "STREAM_PROTOCOL"),
+    # options
+    "OptionIssue": ("jarvis_hep_sampler.options", "OptionIssue"),
+    "ReservedSettings": ("jarvis_hep_sampler.options", "ReservedSettings"),
+    "parse_reserved": ("jarvis_hep_sampler.options", "parse_reserved"),
+    "schema_issues": ("jarvis_hep_sampler.options", "schema_issues"),
+    "split_bounds": ("jarvis_hep_sampler.options", "split_bounds"),
+    # loading and running
+    "LoadedPlugin": ("jarvis_hep_sampler.loader", "LoadedPlugin"),
+    "load_plugin_file": ("jarvis_hep_sampler.loader", "load_plugin_file"),
+    "load_plugin_source": ("jarvis_hep_sampler.loader", "load_plugin_source"),
+    "IndependentSession": ("jarvis_hep_sampler.session", "IndependentSession"),
+    # registry and template
+    "LoadResult": ("jarvis_hep_sampler.registry", "LoadResult"),
+    "RegistryEntry": ("jarvis_hep_sampler.registry", "RegistryEntry"),
+    "SamplerRegistry": ("jarvis_hep_sampler.registry", "SamplerRegistry"),
+    "samplers_home": ("jarvis_hep_sampler.registry", "samplers_home"),
+    "render_template": ("jarvis_hep_sampler.template", "render_template"),
+    "write_template": ("jarvis_hep_sampler.template", "write_template"),
+}
+
+# Kept literal so linters and IDEs see the exports; a test pins it to _LAZY_ATTRS.
+__all__ = [
+    "API_VERSION",
+    "BlockContext",
+    "DrawContext",
+    "IndependentSession",
+    "LoadResult",
+    "LoadedPlugin",
+    "LogBackend",
+    "OptionIssue",
+    "PluginConfigError",
+    "PluginError",
+    "PluginLoadError",
+    "PluginMetadata",
+    "PluginOutputError",
+    "PluginResumeMismatch",
+    "RegistryEntry",
+    "RegistryError",
+    "ReservedSettings",
+    "STREAM_PROTOCOL",
+    "SamplerLog",
+    "SamplerRegistry",
+    "SamplerStats",
+    "SetupContext",
+    "__version__",
+    "check_method_name",
+    "load_plugin_file",
+    "load_plugin_source",
+    "parse_reserved",
+    "render_template",
+    "samplers_home",
+    "schema_issues",
+    "split_bounds",
+    "write_template",
+]
+
+if TYPE_CHECKING:
+    from jarvis_hep_sampler.contexts import BlockContext, DrawContext, SamplerStats, SetupContext
+    from jarvis_hep_sampler.errors import (
+        PluginConfigError,
+        PluginError,
+        PluginLoadError,
+        PluginOutputError,
+        PluginResumeMismatch,
+        RegistryError,
+    )
+    from jarvis_hep_sampler.loader import LoadedPlugin, load_plugin_file, load_plugin_source
+    from jarvis_hep_sampler.log import LogBackend, SamplerLog
+    from jarvis_hep_sampler.metadata import API_VERSION, PluginMetadata, check_method_name
+    from jarvis_hep_sampler.options import (
+        OptionIssue,
+        ReservedSettings,
+        parse_reserved,
+        schema_issues,
+        split_bounds,
+    )
+    from jarvis_hep_sampler.registry import (
+        LoadResult,
+        RegistryEntry,
+        SamplerRegistry,
+        samplers_home,
+    )
+    from jarvis_hep_sampler.session import IndependentSession
+    from jarvis_hep_sampler.streams import STREAM_PROTOCOL
+    from jarvis_hep_sampler.template import render_template, write_template
+
+
+def __getattr__(name: str) -> Any:
+    target = _LAZY_ATTRS.get(name)
+    if target is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attr = target
+    value = getattr(import_module(module_name), attr)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(__all__) | set(globals()))

jarvis_hep_sampler/_check_restart.py

@@ -0,0 +1,29 @@
+"""Private subprocess worker used by testing.check_plugin; no public CLI."""
+
+from __future__ import annotations
+
+import base64
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+from jarvis_hep_sampler.loader import load_plugin_source
+from jarvis_hep_sampler.session import IndependentSession
+
+
+def main() -> None:
+    data = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8"))
+    plugin = load_plugin_source(base64.b64decode(data["source"]), origin=data["origin"])
+    session = IndependentSession(
+        plugin, seed=data["seed"], point_number=data["point_number"],
+        variable_names=data["variable_names"], options=data["options"],
+    )
+    session.prepare()
+    tail = np.array([u for _, u in session.points(data["start"])])
+    np.save(sys.argv[2], tail, allow_pickle=False)
+
+
+if __name__ == "__main__":
+    main()

jarvis_hep_sampler/_version.py

@@ -0,0 +1 @@
+__version__ = "1.0.0"

jarvis_hep_sampler/contexts.py

@@ -0,0 +1,69 @@
+"""Read-only contexts passed to plug-in functions (technical design §4.3)."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from jarvis_hep_sampler.log import SamplerLog
+from jarvis_hep_sampler.streams import prepare_stream
+
+
+@dataclass(frozen=True)
+class SetupContext:
+    """Argument of ``prepare(setup)`` and of the settings hook."""
+
+    seed: int
+    point_number: int
+    dimension: int
+    variable_names: tuple[str, ...]
+    options: Mapping[str, Any]
+    log: SamplerLog
+
+    def rng(self) -> np.random.Generator:
+        """A fresh generator on the stream reserved for ``prepare``."""
+        return prepare_stream(self.seed)
+
+
+@dataclass(frozen=True)
+class DrawContext:
+    """First argument of ``draw(context, rng, plan)``."""
+
+    candidate_index: int
+    point_number: int
+    dimension: int
+    variable_names: tuple[str, ...]
+    options: Mapping[str, Any]
+    log: SamplerLog
+
+
+@dataclass(frozen=True)
+class BlockContext:
+    """First argument of ``draw_block(context, rng, plan)``.
+
+    The block covers candidate indices
+    ``[block_index * block_size, (block_index + 1) * block_size)``.
+    """
+
+    block_index: int
+    block_size: int
+    point_number: int
+    dimension: int
+    variable_names: tuple[str, ...]
+    options: Mapping[str, Any]
+    log: SamplerLog
+
+
+@dataclass(frozen=True)
+class SamplerStats:
+    """Read-only counters passed to the summary, stop-reason and status hooks."""
+
+    candidates: int
+    accepted: int
+    point_number: int
+
+
+__all__ = ["BlockContext", "DrawContext", "SamplerStats", "SetupContext"]

jarvis_hep_sampler/errors.py

@@ -0,0 +1,81 @@
+"""Exceptions raised for sampler plug-ins (technical design §4.7).
+
+Every error carries the context Jarvis-HEP prints in the sampler failure
+record: the method, the plug-in source, the task-card path and, for output
+errors, the candidate index. The plug-in's own exception, when there is one,
+is chained as ``__cause__`` so its traceback reaches the log file.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class PluginError(Exception):
+    """Base class for all plug-in errors."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        method: str | None = None,
+        source: str | None = None,
+        yaml_path: str | None = None,
+        candidate_index: int | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.message = message
+        self.method = method
+        self.source = source
+        self.yaml_path = yaml_path
+        self.candidate_index = candidate_index
+
+    def context_rows(self) -> list[tuple[str, Any]]:
+        """Return the known context as ``(field, value)`` rows."""
+        rows: list[tuple[str, Any]] = []
+        if self.method:
+            rows.append(("method", self.method))
+        if self.candidate_index is not None:
+            rows.append(("candidate", self.candidate_index))
+        if self.yaml_path:
+            rows.append(("at", self.yaml_path))
+        if self.source:
+            rows.append(("source", self.source))
+        return rows
+
+    def __str__(self) -> str:
+        rows = self.context_rows()
+        if not rows:
+            return self.message
+        details = ", ".join(f"{key} {value}" for key, value in rows)
+        return f"{self.message} ({details})"
+
+
+class PluginLoadError(PluginError):
+    """The plug-in file cannot be imported or does not follow the file convention."""
+
+
+class PluginConfigError(PluginError):
+    """The options, ``validate_options`` or ``prepare`` rejected the configuration."""
+
+
+class PluginOutputError(PluginError):
+    """``draw`` / ``draw_block`` raised, or returned a point that breaks the output contract."""
+
+
+class PluginResumeMismatch(PluginError):
+    """A resumed scan does not match the identity recorded by the original run."""
+
+
+class RegistryError(PluginError):
+    """The user sampler registry is unreadable or one of its entries is broken."""
+
+
+__all__ = [
+    "PluginConfigError",
+    "PluginError",
+    "PluginLoadError",
+    "PluginOutputError",
+    "PluginResumeMismatch",
+    "RegistryError",
+]

jarvis_hep_sampler/loader.py

@@ -0,0 +1,237 @@
+"""Load a plug-in file (technical design §4.1).
+
+The file is read once, hashed, compiled and executed from the same bytes, so
+the recorded SHA-256 always matches the code that runs. ``sys.path`` is not
+touched; the module is registered under a private name derived from its path
+and content.
+
+Importing a plug-in runs user code. Only ``Jarvis sampler load`` / ``check``
+and the run preflight do it; ``Jarvis validate`` / ``man`` read the registry.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import inspect
+import sys
+import types
+from collections.abc import Callable, Iterable, Mapping
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from jarvis_hep_sampler.errors import PluginError, PluginLoadError
+from jarvis_hep_sampler.metadata import (
+    GENERATIONAL,
+    INDEPENDENT,
+    SUPPORTED_PROFILES,
+    PluginMetadata,
+    parse_metadata,
+)
+from jarvis_hep_sampler.options import check_options_schema
+
+# name -> number of positional arguments the function must accept.
+PROFILE_FUNCTIONS: Mapping[str, Mapping[str, int]] = {
+    INDEPENDENT: {"prepare": 1, "draw": 3, "draw_block": 3},
+    GENERATIONAL: {"init_state": 1, "ask": 3, "tell": 3},
+}
+OPTIONAL_HOOKS: Mapping[str, int] = {
+    "validate_options": 1,
+    "settings_rows": 2,
+    "summary_rows": 2,
+    "stop_reason": 2,
+    "status": 2,
+}
+_ARGUMENTS = {
+    "prepare": "(setup)",
+    "draw": "(context, rng, plan)",
+    "draw_block": "(context, rng, plan)",
+    "init_state": "(setup)",
+    "ask": "(state, context, rng)",
+    "tell": "(state, context, results)",
+    "validate_options": "(options)",
+    "settings_rows": "(setup, plan)",
+    "summary_rows": "(plan, stats)",
+    "stop_reason": "(plan, stats)",
+    "status": "(plan, stats)",
+}
+
+
+def sha256_bytes(payload: bytes) -> str:
+    return hashlib.sha256(payload).hexdigest()
+
+
+@dataclass(frozen=True)
+class LoadedPlugin:
+    """An imported, structurally valid plug-in."""
+
+    metadata: PluginMetadata
+    options_schema: Mapping[str, Any] | None
+    functions: Mapping[str, Callable[..., Any]]
+    source_path: Path
+    source_sha256: str
+    source_bytes: bytes = field(repr=False)
+    module: types.ModuleType = field(repr=False)
+
+    @property
+    def method(self) -> str:
+        return self.metadata.method
+
+    @property
+    def hooks(self) -> tuple[str, ...]:
+        """Names of the functions the plug-in defines, sorted."""
+        return tuple(sorted(self.functions))
+
+    def function(self, name: str) -> Callable[..., Any] | None:
+        return self.functions.get(name)
+
+
+def load_plugin_file(path: str | Path, *, reserved_names: Iterable[str] = ()) -> LoadedPlugin:
+    """Import the plug-in at *path* and check it against the file convention."""
+    source = Path(path).expanduser()
+    if not source.is_file():
+        raise PluginLoadError(f"plug-in file not found: {source}", source=str(source))
+    resolved = source.resolve()
+    try:
+        payload = resolved.read_bytes()
+    except OSError as exc:
+        raise PluginLoadError(f"cannot read {resolved}: {exc}", source=str(resolved)) from exc
+    return load_plugin_source(payload, origin=resolved, reserved_names=reserved_names)
+
+
+def load_plugin_source(
+    payload: bytes, *, origin: str | Path, reserved_names: Iterable[str] = ()
+) -> LoadedPlugin:
+    """Import a plug-in from its source bytes; *origin* names it in tracebacks."""
+    origin_path = Path(origin)
+    origin_text = str(origin_path)
+    digest = sha256_bytes(payload)
+    module_name = (
+        f"_jarvis_sampler_plugin_{sha256_bytes(origin_text.encode('utf-8'))[:12]}_{digest[:12]}"
+    )
+    try:
+        code = compile(payload, origin_text, "exec", dont_inherit=True)
+    except SyntaxError as exc:
+        raise PluginLoadError(
+            f"cannot compile the plug-in: {exc.msg} (line {exc.lineno})", source=origin_text
+        ) from exc
+
+    module = types.ModuleType(module_name)
+    module.__file__ = origin_text
+    previous = sys.modules.get(module_name)
+    sys.modules[module_name] = module
+    try:
+        exec(code, module.__dict__)
+        plugin = _inspect_module(
+            module,
+            origin=origin_path,
+            digest=digest,
+            payload=payload,
+            reserved_names=reserved_names,
+        )
+    except PluginError:
+        _restore_module(module_name, previous)
+        raise
+    except Exception as exc:
+        _restore_module(module_name, previous)
+        raise PluginLoadError(
+            f"importing the plug-in failed: {type(exc).__name__}: {exc}", source=origin_text
+        ) from exc
+    except BaseException:
+        _restore_module(module_name, previous)
+        raise
+    return plugin
+
+
+def _restore_module(name: str, previous: types.ModuleType | None) -> None:
+    if previous is None:
+        sys.modules.pop(name, None)
+    else:
+        sys.modules[name] = previous
+
+
+def _inspect_module(
+    module: types.ModuleType,
+    *,
+    origin: Path,
+    digest: str,
+    payload: bytes,
+    reserved_names: Iterable[str],
+) -> LoadedPlugin:
+    namespace = module.__dict__
+    source = str(origin)
+    if "PLUGIN" not in namespace:
+        raise PluginLoadError("the plug-in defines no PLUGIN dict", source=source)
+    try:
+        metadata = parse_metadata(namespace["PLUGIN"], reserved_names=reserved_names)
+    except PluginLoadError as exc:
+        exc.source = exc.source or source
+        raise
+    method = metadata.method
+
+    def fail(message: str) -> PluginLoadError:
+        return PluginLoadError(message, method=method, source=source)
+
+    if metadata.profile not in SUPPORTED_PROFILES:
+        raise fail(
+            f"profile {metadata.profile!r} is reserved for a later Jarvis-HEP-Sampler; "
+            f"this version runs: {', '.join(SUPPORTED_PROFILES)}"
+        )
+
+    schema = namespace.get("OPTIONS_SCHEMA")
+    if schema is not None:
+        try:
+            schema = check_options_schema(schema)
+        except PluginLoadError as exc:
+            raise fail(exc.message) from exc
+
+    wanted = {**PROFILE_FUNCTIONS[metadata.profile], **OPTIONAL_HOOKS}
+    functions: dict[str, Callable[..., Any]] = {}
+    for name, arity in wanted.items():
+        if name not in namespace:
+            continue
+        candidate = namespace[name]
+        if not callable(candidate):
+            raise fail(f"{name} must be a function {_ARGUMENTS[name]}")
+        _check_arity(name, candidate, arity, fail)
+        functions[name] = candidate
+
+    if metadata.profile == INDEPENDENT:
+        if metadata.block_size is not None:
+            if "draw_block" not in functions:
+                raise fail("PLUGIN['block_size'] is set, so the plug-in must define draw_block()")
+        elif "draw" not in functions:
+            if "draw_block" in functions:
+                raise fail("draw_block() needs PLUGIN['block_size']; or define draw() instead")
+            raise fail("the independent profile needs a draw(context, rng, plan) function")
+
+    return LoadedPlugin(
+        metadata=metadata,
+        options_schema=schema,
+        functions=functions,
+        source_path=origin,
+        source_sha256=digest,
+        source_bytes=payload,
+        module=module,
+    )
+
+
+def _check_arity(name: str, function: Callable[..., Any], arity: int, fail) -> None:
+    try:
+        signature = inspect.signature(function)
+    except (TypeError, ValueError):
+        return
+    try:
+        signature.bind(*([None] * arity))
+    except TypeError as exc:
+        raise fail(f"{name} must accept {_ARGUMENTS[name]}: {exc}") from exc
+
+
+__all__ = [
+    "LoadedPlugin",
+    "OPTIONAL_HOOKS",
+    "PROFILE_FUNCTIONS",
+    "load_plugin_file",
+    "load_plugin_source",
+    "sha256_bytes",
+]

jarvis_hep_sampler/log.py

@@ -0,0 +1,97 @@
+"""Logging facade handed to plug-ins (technical design §4.5).
+
+Plug-ins never touch a logger directly. They call the ``log`` object they get
+in every context. Inside Jarvis-HEP the backend writes to
+``Jarvis-HEP.Sampler.<Method>`` in the Jarvis layout; offline (``check``,
+:mod:`jarvis_hep_sampler.testing`) it writes to the standard ``logging``
+logger ``jarvis_hep_sampler.<Method>``.
+
+The runtime alone writes the lifecycle records (start, settings, ready,
+checkpoint, failure, result, summary). A plug-in writes only step tables (S5),
+adjustments (S7) and saved-file lines (S11), plus free-form messages.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterable, Sequence
+from typing import Any, Protocol
+
+
+class LogBackend(Protocol):
+    """Where :class:`SamplerLog` sends its records."""
+
+    def emit(self, level: int, message: str) -> None: ...
+
+    def emit_table(self, level: int, title: str, rows: Sequence[tuple[str, Any]]) -> None: ...
+
+
+def format_table(title: str, rows: Iterable[tuple[str, Any]]) -> str:
+    """Render a two-column ``Field | Value`` table like Jarvis-HEP's sampler log."""
+    normalized = [(str(key), str(value)) for key, value in rows]
+    if not normalized:
+        return str(title)
+    key_width = max(len("Field"), *(len(key) for key, _ in normalized))
+    value_width = max(len("Value"), *(len(value) for _, value in normalized))
+    lines = [
+        str(title),
+        f"{'Field':<{key_width}} | {'Value':<{value_width}}",
+        f"{'-' * key_width} | {'-' * value_width}",
+    ]
+    lines.extend(f"{key:<{key_width}} | {value}" for key, value in normalized)
+    return "\n".join(lines)
+
+
+class StdlibLogBackend:
+    """Offline backend: the standard ``logging`` logger ``jarvis_hep_sampler.<method>``."""
+
+    def __init__(self, method: str) -> None:
+        self.logger = logging.getLogger(f"jarvis_hep_sampler.{method}")
+
+    def emit(self, level: int, message: str) -> None:
+        self.logger.log(level, message)
+
+    def emit_table(self, level: int, title: str, rows: Sequence[tuple[str, Any]]) -> None:
+        self.logger.log(level, format_table(title, rows))
+
+
+class SamplerLog:
+    """The ``log`` object a plug-in receives in ``SetupContext`` / ``DrawContext``."""
+
+    def __init__(self, method: str, backend: LogBackend | None = None) -> None:
+        self._method = str(method)
+        self._backend: LogBackend = backend if backend is not None else StdlibLogBackend(method)
+
+    @property
+    def method(self) -> str:
+        return self._method
+
+    def step(self, title: str, rows: Iterable[tuple[str, Any]] = ()) -> None:
+        """S5: one table per generation or stage, e.g. ``log.step("Chains ->", rows)``."""
+        self._backend.emit_table(
+            logging.INFO, str(title), [(str(key), value) for key, value in rows]
+        )
+
+    def adjust(self, name: str, old: Any, new: Any, why: str) -> None:
+        """S7: a setting that was changed, ignored or clamped."""
+        self._backend.emit(logging.WARNING, f"{name}: {old} -> {new} ({why})")
+
+    def saved(self, what: str, path: Any) -> None:
+        """S11: a result file the plug-in wrote."""
+        self._backend.emit(logging.WARNING, f"{what} saved to {path}")
+
+    def debug(self, message: str, *args: Any) -> None:
+        self._backend.emit(logging.DEBUG, _format(message, args))
+
+    def info(self, message: str, *args: Any) -> None:
+        self._backend.emit(logging.INFO, _format(message, args))
+
+    def warning(self, message: str, *args: Any) -> None:
+        self._backend.emit(logging.WARNING, _format(message, args))
+
+
+def _format(message: str, args: tuple[Any, ...]) -> str:
+    return str(message) % args if args else str(message)
+
+
+__all__ = ["LogBackend", "SamplerLog", "StdlibLogBackend", "format_table"]