erislib 0.1.0__py3-none-any.whl

eris/__init__.py

@@ -0,0 +1,43 @@
+"""Eris public API."""
+
+from __future__ import annotations
+
+from eris._version import __version__
+from eris.api import embed, encode, pool
+from eris.cli._config_file import ConfigIssue
+from eris.cli.validate import validate_config
+from eris.discovery import (
+    get_embedder_info,
+    get_representation_info,
+    list_embedders,
+    list_representations,
+)
+from eris.exceptions import RepresentationNotFoundError
+from eris.export.pyg import to_pyg
+from eris.export.table import feature_table
+from eris.export.torch import to_torch
+from eris.pipelines.batch_encode import batch_embed, batch_encode
+from eris.pipelines.pack import pack_artifacts as pack
+
+__all__ = [
+    "ConfigIssue",
+    "RepresentationNotFoundError",
+    "__version__",
+    "batch_embed",
+    "batch_encode",
+    "embed",
+    "encode",
+    "feature_table",
+    "get_embedder_info",
+    "get_representation_info",
+    "list_embedders",
+    "list_representations",
+    "pack",
+    "pool",
+    "to_pyg",
+    "to_torch",
+    "validate_config",
+]
+
+# register built-in encoders
+from eris import encoders as _encoders  # noqa: F401

eris/__main__.py

@@ -0,0 +1,5 @@
+"""Entry point for ``python -m eris``."""
+
+from eris.cli.main import main
+
+raise SystemExit(main())

eris/_provenance.py

@@ -0,0 +1,104 @@
+"""Provenance sidecar (``provenance/1``) shared by the ecosystem CLIs.
+
+Copied into each library rather than imported, so the libraries stay independent.
+Keep copies in sync with the contract's ``provenance.schema.json``.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import platform
+from datetime import UTC, datetime
+from importlib import metadata
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Mapping
+    from pathlib import Path
+
+SCHEMA_VERSION = "provenance/1"
+_CHUNK = 1 << 20
+
+
+def utc_now() -> str:
+    """Return the current UTC time as ``YYYY-MM-DDTHH:MM:SSZ``."""
+    return datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+def sha256_path(path: Path) -> str:
+    r"""Hash a file, or a directory as the sorted ``relative/path\\tfile_sha256`` lines of its files."""
+    if path.is_dir():
+        lines = (
+            f"{p.relative_to(path).as_posix()}\t{_sha256_file(p)}\n"
+            for p in sorted(path.rglob("*"))
+            if p.is_file()
+        )
+        return hashlib.sha256("".join(lines).encode()).hexdigest()
+    return _sha256_file(path)
+
+
+def _sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(_CHUNK), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def canonical_json(value: Any) -> str:
+    """Serialize deterministically; paths and other non-JSON scalars become strings."""
+    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str)
+
+
+def fingerprint(parameters: Mapping[str, Any]) -> str:
+    """Return ``sha256:<hex>`` of the canonical JSON of ``parameters``."""
+    return "sha256:" + hashlib.sha256(canonical_json(parameters).encode()).hexdigest()
+
+
+def _package_version(name: str) -> str | None:
+    try:
+        return metadata.version(name)
+    except metadata.PackageNotFoundError:
+        return None
+
+
+def write_provenance(
+    output: Path,
+    *,
+    tool: str,
+    version: str,
+    command: str,
+    created_at: str,
+    parameters: Mapping[str, Any],
+    inputs: Iterable[Path] = (),
+    packages: Iterable[str] = (),
+    output_info: Mapping[str, int] | None = None,
+    extensions: Mapping[str, Any] | None = None,
+) -> Path:
+    """Write ``<output>.provenance.json`` next to ``output`` (a file or a directory) and return its path.
+
+    Call it after ``output`` is complete; ``created_at`` is the ``utc_now()`` taken when the run started.
+    """
+    document = {
+        "schema_version": SCHEMA_VERSION,
+        "tool": {"name": tool, "version": version},
+        "command": command,
+        "created_at": created_at,
+        "finished_at": utc_now(),
+        "parameters": json.loads(canonical_json(parameters)),
+        "config_fingerprint": fingerprint(parameters),
+        "inputs": [{"path": str(path), "sha256": sha256_path(path)} for path in inputs],
+        "outputs": [{"path": str(output), "sha256": sha256_path(output), **(output_info or {})}],
+        "environment": {
+            "python": platform.python_version(),
+            "platform": platform.platform(),
+            "packages": {name: _package_version(name) for name in packages},
+        },
+        "extensions": dict(extensions or {}),
+    }
+    target = output.with_name(output.name + ".provenance.json")
+    target.write_text(
+        json.dumps(document, indent=2, ensure_ascii=False, default=str) + "\n", encoding="utf-8"
+    )
+    return target

eris/_version.py

@@ -0,0 +1,3 @@
+"""Package version."""
+
+__version__ = "0.1.0"

eris/api.py

@@ -0,0 +1,203 @@
+"""Public encode and embed entry points."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+from eris.core.artifacts import StructureArtifact
+from eris.core.config import CanonicalRepresentationConfig
+from eris.core.registry import (
+    get_encoder,
+    resolve_representation_config,
+    validate_resolved_representation_config,
+)
+from eris.core.types import StructureInput
+from eris.exceptions import ConfigError, RepresentationNotFoundError
+from eris.pooling import pool
+
+if TYPE_CHECKING:
+    from eris.core.learned.config import OutputMode
+    from eris.core.learned.execution import Precision
+
+
+class _NoOverride:
+    """Sentinel type for an argument that was not overridden."""
+
+    def __repr__(self) -> str:
+        return "<no override>"
+
+
+_NO_OVERRIDE = _NoOverride()
+
+
+def _normalize_chain_id(value: object) -> tuple[str, ...] | None:
+    if value is None:
+        return None
+    if isinstance(value, str):
+        normalized = (value,)
+    elif isinstance(value, Sequence) and not isinstance(value, (bytes, bytearray)):
+        normalized = tuple(value)
+    else:
+        raise ConfigError("chain_id must be None, a string, or a sequence of strings")
+    if not normalized:
+        raise ConfigError("chain_id sequence must not be empty")
+    if any(not isinstance(chain, str) for chain in normalized):
+        raise ConfigError("chain_id sequence entries must be strings")
+    if len(normalized) != len(set(normalized)):
+        raise ConfigError("chain_id sequence must not contain duplicates")
+    return normalized
+
+
+def _validate_model_index(value: object) -> None:
+    if type(value) is not int or value < 0:
+        raise ConfigError("model_index must be a plain integer >= 0")
+
+
+def _coerce_structure_input(
+    input: str | Path | StructureInput,  # noqa: A002
+    *,
+    chain_id: str | Sequence[str] | _NoOverride | None,
+    model_index: int | _NoOverride,
+) -> StructureInput:
+    if isinstance(input, StructureInput):
+        if chain_id is not _NO_OVERRIDE:
+            normalized_chain_id = _normalize_chain_id(chain_id)
+            if normalized_chain_id != input.chain_id:
+                raise ConfigError(
+                    f"chain_id override conflicts with StructureInput: {chain_id!r} != {input.chain_id!r}"
+                )
+        if model_index is not _NO_OVERRIDE:
+            _validate_model_index(model_index)
+            if model_index != input.model_index:
+                raise ConfigError(
+                    "model_index override conflicts with StructureInput: "
+                    f"{model_index!r} != {input.model_index!r}"
+                )
+        return input
+
+    if not isinstance(input, (str, Path)):
+        raise TypeError("input must be a string, pathlib.Path, or StructureInput")
+
+    effective_chain = None if chain_id is _NO_OVERRIDE else chain_id
+    effective_model = 0 if model_index is _NO_OVERRIDE else model_index
+    effective_chain = _normalize_chain_id(effective_chain)
+    _validate_model_index(effective_model)
+    return StructureInput(
+        source=str(input),
+        chain_id=effective_chain,
+        model_index=cast("int", effective_model),
+    )
+
+
+def _resolve_config(
+    representation: str,
+    *,
+    config: CanonicalRepresentationConfig | None,
+    params: dict[str, Any],
+) -> CanonicalRepresentationConfig:
+    if config is not None:
+        return validate_resolved_representation_config(
+            representation,
+            config=config,
+            params=params,
+        )
+    return resolve_representation_config(
+        representation,
+        params=params,
+    )
+
+
+def encode(
+    input: str | Path | StructureInput,  # noqa: A002
+    *,
+    representation: str,
+    chain_id: str | Sequence[str] | _NoOverride | None = _NO_OVERRIDE,
+    model_index: int | _NoOverride = _NO_OVERRIDE,
+    config: CanonicalRepresentationConfig | None = None,
+    output: str | Path | None = None,
+    use_cache: bool = True,
+    **params: Any,
+) -> StructureArtifact:
+    """Encode one structure with a canonical explicit representation."""
+    if type(use_cache) is not bool:
+        raise ConfigError("use_cache must be a bool")
+    if output is not None and not isinstance(output, (str, Path)):
+        raise TypeError("output must be a string, pathlib.Path, or None")
+
+    resolved_config = _resolve_config(
+        representation,
+        config=config,
+        params=params,
+    )
+    structure_input = _coerce_structure_input(
+        input,
+        chain_id=chain_id,
+        model_index=model_index,
+    )
+    encoder = get_encoder(
+        representation,
+        config=resolved_config,
+        use_cache=use_cache,
+    )
+    artifact = encoder.encode(structure_input)
+    if output is not None:
+        artifact.save(Path(output))
+    return artifact
+
+
+def embed(
+    input: str | Path | StructureInput,  # noqa: A002
+    *,
+    representation: str,
+    chain_id: str | Sequence[str] | _NoOverride | None = _NO_OVERRIDE,
+    model_index: int | _NoOverride = _NO_OVERRIDE,
+    output_mode: str = "both",
+    checkpoint_path: str | Path | None = None,
+    device: str = "auto",
+    precision: str = "fp32",
+    output: str | Path | None = None,
+    use_cache: bool = True,
+) -> StructureArtifact:
+    """Embed one structure with a registered learned representation."""
+    from eris.core.learned import StructuralEmbedderRunner, StructuralExecutionConfig  # noqa: PLC0415
+    from eris.core.learned.registry import learned_registry  # noqa: PLC0415
+
+    try:
+        info = learned_registry.get_info(representation)
+    except KeyError as error:
+        raise RepresentationNotFoundError(f"Unknown embedder {representation!r}") from error
+
+    if type(use_cache) is not bool:
+        raise ConfigError("use_cache must be a bool")
+    if output is not None and not isinstance(output, (str, Path)):
+        raise TypeError("output must be a string, pathlib.Path, or None")
+    if checkpoint_path is not None and not isinstance(checkpoint_path, (str, Path)):
+        raise TypeError("checkpoint_path must be a string, pathlib.Path, or None")
+
+    structure_input = _coerce_structure_input(
+        input,
+        chain_id=chain_id,
+        model_index=model_index,
+    )
+    config = info.config_type(
+        representation_name=info.representation_name,
+        output_mode=cast("OutputMode", output_mode),
+    )
+    execution = StructuralExecutionConfig(device=device, precision=cast("Precision", precision))
+
+    artifact = StructuralEmbedderRunner().run(
+        structure_input,
+        config=config,
+        model=info.model_descriptor,
+        checkpoint_path=checkpoint_path,
+        execution=execution,
+        use_artifact_cache=use_cache,
+    )
+    if output is not None:
+        artifact.save(Path(output))
+    return artifact
+
+
+__all__ = ["embed", "encode", "pool"]

eris/cli/__init__.py

@@ -0,0 +1 @@
+"""Command-line interface."""

eris/cli/_config_file.py

@@ -0,0 +1,83 @@
+"""``--config`` files whose keys mirror a Typer command's options.
+
+Keys are option parameter names (``--max-length`` -> ``max_length``). Values are checked with the
+option's own Click type (through Typer); the callers install accepted files as the context's ``default_map``,
+so the precedence is defaults < config file < options on the command line.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from typing import TYPE_CHECKING, Any
+
+import typer
+import yaml
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+@dataclass(frozen=True)
+class ConfigIssue:
+    """One problem found in a config file; ``path`` is ``""`` for the root."""
+
+    path: str
+    message: str
+
+
+def read_config(path: Path) -> tuple[dict[str, Any], list[ConfigIssue]]:
+    """Parse a YAML (or JSON) config whose root must be a mapping. An empty file is ``{}``."""
+    try:
+        raw = yaml.safe_load(path.read_text(encoding="utf-8"))
+    except (OSError, yaml.YAMLError) as exc:
+        return {}, [ConfigIssue("", f"cannot read config: {exc}")]
+    if raw is None:
+        return {}, []
+    if not isinstance(raw, dict):
+        return {}, [ConfigIssue("", f"config root must be a mapping, got {type(raw).__name__}")]
+    return raw, []
+
+
+def check_config(command: Any, raw: dict[str, Any], *, exclude: frozenset[str]) -> list[ConfigIssue]:
+    """Check every key and value against ``command``'s options.
+
+    ``exclude`` names options that must stay on the command line.
+    """
+    ctx = typer.Context(command)
+    # duck-typed: Typer >=0.27 vendors Click, so isinstance checks against `click` fail
+    options = {
+        p.name: p for p in command.params if getattr(p, "param_type_name", None) == "option" and p.name
+    }
+    issues: list[ConfigIssue] = []
+    for key, value in raw.items():
+        param = options.get(key)
+        if param is None:
+            issues.append(ConfigIssue(str(key), "unknown key"))
+        elif key in exclude:
+            issues.append(ConfigIssue(key, f"pass {param.opts[0]} on the command line, not in the config"))
+        elif param.multiple and not isinstance(value, list):
+            issues.append(ConfigIssue(key, f"expected a list, got {type(value).__name__}"))
+        # Click would coerce a number into a text or path option
+        elif param.type.name in ("str", "path") and not param.multiple and not isinstance(value, str):
+            issues.append(ConfigIssue(key, f"expected a string, got {type(value).__name__}"))
+        else:
+            try:
+                param.type_cast_value(ctx, value)
+            except typer.BadParameter as exc:
+                issues.append(ConfigIssue(key, exc.format_message()))
+            except AttributeError:  # Click's bool type calls .strip() on non-strings such as 3
+                issues.append(ConfigIssue(key, f"invalid value {value!r}"))
+    return issues
+
+
+def report_issues(issues: list[ConfigIssue], *, as_json: bool) -> int:
+    """Print issues in the contract's format and return the exit code (0 valid, 2 invalid)."""
+    if as_json:
+        typer.echo(json.dumps([asdict(issue) for issue in issues]))
+    else:
+        for issue in issues:
+            typer.echo(f"{issue.path or '<root>'}: {issue.message}", err=True)
+    if not issues and not as_json:
+        typer.echo("config is valid", err=True)
+    return 2 if issues else 0

eris/cli/_shared.py

@@ -0,0 +1,119 @@
+"""Shared CLI pieces: ``--config`` handling, config validation and the provenance sidecar."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+from dataclasses import fields
+from pathlib import Path
+from typing import Any
+
+import typer
+
+import eris
+from eris._provenance import write_provenance
+from eris.cli._config_file import ConfigIssue, check_config, read_config
+from eris.core.registry import get_encoder_info
+
+# Options that stay on the command line and are not provenance parameters.
+CONFIG_EXCLUDE = frozenset({"out", "out_path", "artifact_dir", "results", "checkpoint", "config"})
+
+
+def check_params(representation: str, params: Any) -> list[ConfigIssue]:
+    """Check representation parameters against the representation's own config dataclass."""
+    try:
+        config_type = get_encoder_info(representation).config_type
+    except KeyError:
+        return [ConfigIssue("representation", f"unknown representation {representation!r}")]
+    if not isinstance(params, Mapping):
+        return [ConfigIssue("params", f"expected a mapping, got {type(params).__name__}")]
+    known = {f.name for f in fields(config_type)}
+    issues = [ConfigIssue(f"params.{key}", "unknown parameter") for key in params if key not in known]
+    for key in params.keys() & known:
+        try:
+            config_type(**{key: params[key]})
+        except Exception as exc:  # noqa: BLE001 - any constructor failure is a config problem
+            issues.append(ConfigIssue(f"params.{key}", str(exc)))
+    if not issues:
+        try:
+            config_type(**params)
+        except Exception as exc:  # noqa: BLE001 - any constructor failure is a config problem
+            issues.append(ConfigIssue("params", str(exc)))
+    return issues
+
+
+def config_issues(command: Any, path: Path) -> tuple[dict[str, Any], list[ConfigIssue]]:
+    """Read ``path`` and collect every problem with its options and nested ``params``."""
+    raw, issues = read_config(path)
+    if issues:
+        return raw, issues
+    return raw, mapping_issues(command, raw)
+
+
+def mapping_issues(command: Any, raw: dict[str, Any]) -> list[ConfigIssue]:
+    """Collect every problem of a parsed config with ``command``'s options and nested ``params``."""
+    rest = {key: value for key, value in raw.items() if key != "params"}
+    issues = check_config(command, rest, exclude=CONFIG_EXCLUDE)
+    if "params" in raw:
+        if any(getattr(p, "name", None) == "params" for p in command.params):
+            if isinstance(raw.get("representation"), str):
+                issues += check_params(raw["representation"], raw["params"])
+            elif not isinstance(raw["params"], Mapping):
+                issues.append(
+                    ConfigIssue("params", f"expected a mapping, got {type(raw['params']).__name__}")
+                )
+        else:
+            issues.append(ConfigIssue("params", "unknown key"))
+    return issues
+
+
+def config_callback(ctx: typer.Context, value: Path | None) -> Path | None:
+    """Eager ``--config`` callback: install the file as defaults, so flags still win."""
+    if value is None:
+        return None
+    raw, issues = config_issues(ctx.command, value)
+    if issues:
+        lines = "\n".join(f"  {issue.path or '<root>'}: {issue.message}" for issue in issues)
+        raise typer.BadParameter(f"{value}:\n{lines}", ctx=ctx, param_hint="'--config'")
+    defaults = dict(raw)
+    if "params" in defaults:  # the option takes a JSON string; the file holds a mapping
+        defaults["params"] = json.dumps(defaults["params"])
+    ctx.default_map = {**(ctx.default_map or {}), **defaults}
+    return value
+
+
+CONFIG_OPTION = typer.Option(
+    "--config",
+    callback=config_callback,
+    is_eager=True,
+    dir_okay=False,
+    help="YAML/JSON file with option values; representation parameters go under 'params'.",
+)
+
+
+def write_run_provenance(
+    ctx: typer.Context,
+    output: Path,
+    *,
+    command: str,
+    created_at: str,
+    inputs: list[Path],
+    extensions: dict[str, Any] | None = None,
+) -> None:
+    """Write the provenance sidecar for a finished run."""
+    parameters = {k: v for k, v in ctx.params.items() if k not in CONFIG_EXCLUDE}
+    for path_option in ("input", "manifest", "folder", "index", "artifact_root"):
+        parameters.pop(path_option, None)
+    if parameters.get("params") is not None:  # record the parsed mapping, not the JSON string
+        parameters["params"] = json.loads(parameters["params"])
+    write_provenance(
+        output,
+        tool="eris",
+        version=eris.__version__,
+        command=command,
+        created_at=created_at,
+        parameters=parameters,
+        inputs=inputs,
+        packages=("erislib", "numpy", "polars"),
+        extensions=extensions,
+    )

eris/cli/_utils.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import NoReturn
+
+import typer
+from platformdirs import user_cache_dir
+
+
+def usage_error(message: str) -> NoReturn:
+    typer.echo(message, err=True)
+    raise typer.Exit(code=2)
+
+
+def parse_json_object(params: str | None) -> dict[str, object]:
+    if params is None:
+        return {}
+    try:
+        value = json.loads(params)
+    except json.JSONDecodeError as error:
+        usage_error(f"Invalid JSON for --params: {error.msg}")
+    if not isinstance(value, dict):
+        usage_error("--params must be a JSON object")
+    return value
+
+
+def cache_root() -> Path:
+    return Path(user_cache_dir("eris"))