diffract-core 0.2.1__py3-none-any.whl

diffract/session/session.py

@@ -0,0 +1,260 @@
+"""Main Session API for the diffract package.
+
+This module implements the primary user interface for neural network parameter
+analysis workflows, providing a high-level API for model management, kernel
+execution, and result retrieval through dependency injection.
+"""
+
+from __future__ import annotations
+
+import logging
+import types
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Self
+
+from dependency_injector.wiring import Provide, inject
+
+from diffract.containers import MainContainer, create_main_container
+from diffract.core.data.nn.aggregates.repository import AggregateRepository
+from diffract.core.data.nn.aggregates.view import AggregateView
+from diffract.core.data.nn.params.interface import IParameterRepository, IParameterView
+from diffract.core.data.nn.params.schema import ParameterType
+from diffract.session.field_cache import SessionFieldCache
+from diffract.session.utils import filter_aggregate_view, filter_parameter_view
+
+logger = logging.getLogger(__name__)
+
+
+class Session:
+    """Main session class for neural network parameter analysis.
+
+    Provides a high-level interface for adding models, applying computational
+    kernels, and retrieving results. Manages the complete lifecycle of parameter
+    analysis workflows through dependency injection and modular components.
+
+    Example:
+        >>> session = Session()
+        >>> session.models.add(my_model, model_id="bert-base")
+        >>> session.compute.apply("frob_norm", "stable_rank")
+        >>> results = session.results.export_metrics(
+        ...     "frob_norm", "stable_rank", export_format="pandas"
+        ... )
+    """
+
+    def __init__(
+        self,
+        profile: str | None = "ram",
+        config_path: str | Path | None = None,
+        container: MainContainer | None = None,
+        other_session: Session | None = None,
+    ) -> None:
+        """Initialize a new session.
+
+        Args:
+            profile: Built-in profile name for quick setup:
+                - "ram": RAM storage, no persistence (fast experiments)
+                - "local": SQLite storage in .diffract/ (persistent, simple)
+                - "hybrid": SQLite + HDF5 (persistent, optimized for large arrays)
+            config_path: Path to configuration file (YAML/JSON/INI). Takes
+                priority over profile if both are provided.
+            container: Pre-configured MainContainer instance. If None,
+                creates new container using config_path or profile.
+            other_session: Another Session instance to copy configuration from.
+        """
+        from diffract.session.namespaces import (
+            ComputeNamespace,
+            ModelsNamespace,
+            ResultsNamespace,
+            UtilsNamespace,
+            VizNamespace,
+        )
+
+        # For namespaces
+        self._field_cache = SessionFieldCache()
+
+        if container is None:
+            self.__container = create_main_container(config_path, profile=profile)
+        else:
+            self.__container = container
+
+        self.__context_depth = 0
+
+        with self:
+            logger.debug("Loading parameters...")
+            self.__init_repos()
+            logger.info(
+                "Session initialized with %d existing parameters, %d aggregates",
+                len(self.__param_repo),
+                len(self.__agg_repo),
+            )
+
+        self.models = ModelsNamespace(self)
+        self.compute = ComputeNamespace(self)
+        self.results = ResultsNamespace(self)
+        self.viz = VizNamespace(self)
+        self.utils = UtilsNamespace(self)
+
+        if other_session is not None:
+            self.utils.merge_other_session(other_session, verify=False)
+
+    def __enter__(self) -> Self | nullcontext[Self]:
+        """Enter session context and initialize resources."""
+        self.__context_depth += 1
+        if self.__context_depth == 1:
+            logger.debug("Init resources called")
+            self.__container.init_resources()
+            return self
+        return nullcontext()
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        """Exit session context and shutdown resources."""
+        self.__context_depth -= 1
+        if self.__context_depth == 0:
+            logger.debug("Shutdown resources called")
+            self.__container.shutdown_resources()
+
+    @inject
+    def __init_repos(
+        self,
+        parameter_repository: IParameterRepository = Provide["nn.parameter_repository"],
+        aggregate_repository: AggregateRepository = Provide["nn.aggregate_repository"],
+    ) -> None:
+        """Inject dependencies from the container."""
+        self.__param_repo = parameter_repository
+        self.__agg_repo = aggregate_repository
+
+    def filter(
+        self,
+        param_ids: list[str] | None = None,
+        param_names: list[str] | None = None,
+        param_types: list[ParameterType] | None = None,
+        model_ids: list[str] | None = None,
+    ) -> SessionContext:
+        """Create a filtered view of the session's parameters and aggregates.
+
+        Args:
+            param_ids: Parameter IDs to keep, or None for no ID filtering.
+            param_names: Parameter names to keep, or None for no name filtering.
+            param_types: Parameter types to keep, or None for no type filtering.
+            model_ids: Model IDs to keep, or None for no model filtering.
+
+        Returns:
+            A SessionContext scoped to the filtered parameter and aggregate views.
+        """
+        with self:
+            param_view = self.__param_repo.create_view()
+            filtered_param_view = filter_parameter_view(
+                param_view,
+                parameter_ids=param_ids,
+                parameter_names=param_names,
+                parameter_types=param_types,
+                model_ids=model_ids,
+            )
+
+            agg_view = self.__agg_repo.create_view()
+            filtered_agg_view = filter_aggregate_view(
+                agg_view,
+                parameter_names=param_names,
+                model_ids=model_ids,
+            )
+
+            return SessionContext(self, filtered_param_view, filtered_agg_view)
+
+
+class SessionContext:
+    """Filtered, scoped view over a session's parameters and aggregates.
+
+    Wraps a parent Session together with pre-filtered parameter and aggregate
+    views, exposing the same namespaces (models, compute, results, viz, utils)
+    restricted to the filtered scope. Supports further chained filtering.
+    """
+
+    def __init__(
+        self,
+        session: Session,
+        param_filter_context: IParameterView,
+        agg_filter_context: AggregateView,
+    ) -> None:
+        """Initialize a filtered session context.
+
+        Args:
+            session: The parent Session this context is derived from.
+            param_filter_context: Pre-filtered parameter view for this scope.
+            agg_filter_context: Pre-filtered aggregate view for this scope.
+        """
+        from diffract.session.namespaces import (
+            ComputeNamespace,
+            ModelsNamespace,
+            ResultsNamespace,
+            UtilsNamespace,
+            VizNamespace,
+        )
+
+        self.__session = session
+
+        # For namespaces
+        self._param_filter_context = param_filter_context
+        self._agg_filter_context = agg_filter_context
+        self._field_cache = session._field_cache
+
+        self.models = ModelsNamespace(self)
+        self.compute = ComputeNamespace(self)
+        self.results = ResultsNamespace(self)
+        self.viz = VizNamespace(self)
+        self.utils = UtilsNamespace(self)
+
+    def __enter__(self) -> Self | nullcontext[Self]:
+        """Enter the underlying session context."""
+        return self.__session.__enter__()
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        """Exit the underlying session context."""
+        self.__session.__exit__(exc_type, exc_val, exc_tb)
+
+    def filter(
+        self,
+        param_ids: list[str] | None = None,
+        param_names: list[str] | None = None,
+        param_types: list[ParameterType] | None = None,
+        model_ids: list[str] | None = None,
+    ) -> Self:
+        """Further filter this context's parameters and aggregates.
+
+        Args:
+            param_ids: Parameter IDs to keep, or None for no ID filtering.
+            param_names: Parameter names to keep, or None for no name filtering.
+            param_types: Parameter types to keep, or None for no type filtering.
+            model_ids: Model IDs to keep, or None for no model filtering.
+
+        Returns:
+            A new SessionContext scoped to the further-filtered views.
+        """
+        with self:
+            filtered_param_view = filter_parameter_view(
+                self._param_filter_context,
+                parameter_ids=param_ids,
+                parameter_names=param_names,
+                parameter_types=param_types,
+                model_ids=model_ids,
+            )
+
+            filtered_agg_view = filter_aggregate_view(
+                self._agg_filter_context,
+                parameter_names=param_names,
+                model_ids=model_ids,
+            )
+
+            return SessionContext(
+                self.__session, filtered_param_view, filtered_agg_view
+            )

diffract/session/utils.py

@@ -0,0 +1,81 @@
+from diffract.core.data.nn.aggregates.view import AggregateView
+from diffract.core.data.nn.params import ParameterType
+from diffract.core.data.nn.params.interface import IParameterView
+
+
+def filter_parameter_view(
+    view: IParameterView,
+    *,
+    parameter_ids: list[str] | None = None,
+    parameter_names: list[str] | None = None,
+    parameter_types: list[ParameterType] | None = None,
+    model_ids: list[str] | None = None,
+) -> IParameterView:
+    """Apply the given filters to a parameter view.
+
+    Each filter is only applied when a non-empty value is provided. Scalar
+    values are wrapped into single-element lists before filtering.
+
+    Args:
+        view: The parameter view to filter.
+        parameter_ids: Parameter IDs to select, or None to skip ID filtering.
+        parameter_names: Parameter names to filter by, or None to skip.
+        parameter_types: Parameter types to filter by, or None to skip.
+        model_ids: Model IDs to filter by, or None to skip.
+
+    Returns:
+        The filtered parameter view.
+    """
+    if parameter_names:
+        if isinstance(parameter_names, str):
+            parameter_names = [parameter_names]
+        view = view.filter_by_name(*parameter_names)
+
+    if parameter_types:
+        if isinstance(parameter_types, (str, ParameterType)):
+            parameter_types = [parameter_types]
+        view = view.filter_by_ptype(*parameter_types)
+
+    if model_ids:
+        if isinstance(model_ids, str):
+            model_ids = [model_ids]
+        view = view.filter_by_model_id(*model_ids)
+
+    if parameter_ids:
+        if isinstance(parameter_ids, (str, int)):
+            parameter_ids = [parameter_ids]
+        view = view[parameter_ids]
+
+    return view
+
+
+def filter_aggregate_view(
+    view: AggregateView,
+    *,
+    parameter_names: list[str] | None = None,
+    model_ids: list[str] | None = None,
+) -> AggregateView:
+    """Apply the given filters to an aggregate view.
+
+    Each filter is only applied when a non-empty value is provided. Scalar
+    values are wrapped into single-element lists before filtering.
+
+    Args:
+        view: The aggregate view to filter.
+        parameter_names: Context parameter names to filter by, or None to skip.
+        model_ids: Context model IDs to filter by, or None to skip.
+
+    Returns:
+        The filtered aggregate view.
+    """
+    if parameter_names:
+        if isinstance(parameter_names, str):
+            parameter_names = [parameter_names]
+        view = view.filter_by_context_params(*parameter_names)
+
+    if model_ids:
+        if isinstance(model_ids, str):
+            model_ids = [model_ids]
+        view = view.filter_by_context_models(*model_ids)
+
+    return view

diffract/viz/__init__.py

@@ -0,0 +1,42 @@
+"""Refactored visualization stack for diffract.
+
+Submodules that need plotting backends are imported lazily, so the package
+and its data helpers stay importable without the ``viz`` extra installed.
+"""
+
+from typing import Any
+
+__all__ = [
+    "DARK_THEME",
+    "DEFAULT_THEME",
+    "MINIMAL_THEME",
+    "Plot",
+    "Theme",
+    "apply_theme",
+    "load_theme",
+    "plots",
+    "render",
+    "render_from_config",
+]
+
+_RENDERER_EXPORTS = frozenset({"Plot", "load_theme", "render", "render_from_config"})
+_STYLING_EXPORTS = frozenset(
+    {"DARK_THEME", "DEFAULT_THEME", "MINIMAL_THEME", "Theme", "apply_theme"}
+)
+
+
+def __getattr__(name: str) -> Any:
+    if name in _RENDERER_EXPORTS:
+        from . import renderer
+
+        return getattr(renderer, name)
+    if name in _STYLING_EXPORTS:
+        from . import styling
+
+        return getattr(styling, name)
+    if name == "plots":
+        from . import plots
+
+        return plots
+    msg = f"module {__name__!r} has no attribute {name!r}"
+    raise AttributeError(msg)

diffract/viz/data/__init__.py

@@ -0,0 +1,34 @@
+from .detection import detect_data_shape, detect_data_type, detect_field_meta
+from .extraction import (
+    get_field_data,
+    get_field_value,
+    get_field_values,
+)
+from .filtering import apply_filter
+from .ordering import Ordering, OrderMode, as_is, by_key, custom, lexicographic, numeric
+from .provider import DataProvider
+from .types import DataShape, DataType, Entry, EntryContext, FieldMeta, FieldRef
+
+__all__ = [
+    "DataProvider",
+    "DataShape",
+    "DataType",
+    "Entry",
+    "EntryContext",
+    "FieldMeta",
+    "FieldRef",
+    "OrderMode",
+    "Ordering",
+    "apply_filter",
+    "as_is",
+    "by_key",
+    "custom",
+    "detect_data_shape",
+    "detect_data_type",
+    "detect_field_meta",
+    "get_field_data",
+    "get_field_value",
+    "get_field_values",
+    "lexicographic",
+    "numeric",
+]

diffract/viz/data/detection.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import re
+from typing import Any
+
+import numpy as np
+
+from .types import DataShape, DataType, FieldMeta
+
+RE_NUMBER_PATTERN = re.compile(r"[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?")
+
+
+def is_number_regex(string: str) -> bool:
+    """Return True if the string fully matches a numeric pattern."""
+    return RE_NUMBER_PATTERN.fullmatch(string) is not None
+
+
+def detect_field_meta(values: list[Any]) -> FieldMeta:
+    """Detect DataType and DataShape from values."""
+    data_type = detect_data_type(values)
+    data_shape = detect_data_shape(values)
+    return FieldMeta(data_type, data_shape)
+
+
+def detect_data_type(values: list[Any]) -> DataType:
+    """Detect if values are numeric or categorical."""
+    for v in values:
+        if v is None:
+            continue
+        if isinstance(v, bool):
+            return DataType.CATEGORICAL
+        if isinstance(v, str):
+            if is_number_regex(v):
+                continue
+            return DataType.CATEGORICAL
+        if isinstance(v, (int, float, np.integer, np.floating)):
+            continue
+        if isinstance(v, np.ndarray):
+            if not np.issubdtype(v.dtype, np.number):
+                return DataType.CATEGORICAL
+            continue
+        return DataType.CATEGORICAL
+
+    return DataType.NUMERIC
+
+
+def detect_data_shape(values: list[Any]) -> DataShape:
+    """Detect if field represents scalar or vector per entry."""
+    for v in values:
+        if v is None:
+            continue
+        if isinstance(v, np.ndarray) and v.size > 1:
+            return DataShape.VECTOR
+        if isinstance(v, (list, tuple)) and len(v) > 1:
+            return DataShape.VECTOR
+
+    return DataShape.SCALAR

diffract/viz/data/extraction.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from diffract.core.constants import CONTEXT_SEPARATOR, parse_contextual_field_name
+
+from .detection import detect_field_meta
+from .types import DataShape, DataType, Entry, EntryContext
+
+
+def get_field_value(ctx: EntryContext, field: str) -> Any:
+    """Return the value of a field, resolving contextual field variants.
+
+    Args:
+        ctx: Entry context to look the field up in.
+        field: Field name to resolve.
+
+    Returns:
+        The direct field value if present, otherwise the best-matching
+        contextual field value.
+
+    Raises:
+        ValueError: If no matching field is found in the context.
+    """
+    value = ctx.fields.get(field)
+    if value is not None:
+        return value
+
+    agg_pattern = re.compile(rf"^{re.escape(field)}{re.escape(CONTEXT_SEPARATOR)}.+$")
+
+    candidates: list[tuple[str, Any]] = []
+    for key_, value_ in ctx.fields.items():
+        if agg_pattern.match(key_):
+            candidates.append((key_, value_))
+
+    if not candidates:
+        raise ValueError(f"Field {field} not found in entry context")
+
+    return _select_best_contextual_field(candidates, ctx=ctx)
+
+
+def get_field_values(entries: dict[str, Entry], field: str) -> list[Any]:
+    """Return the field value for each entry.
+
+    Args:
+        entries: Mapping of uid to entry.
+        field: Field name to extract.
+
+    Returns:
+        List of field values, one per entry.
+    """
+    values: list[Any] = []
+
+    for entry in entries.values():
+        value = get_field_value(EntryContext.from_entry(entry), field)
+        values.append(value)
+
+    return values
+
+
+def get_field_data(
+    entries: dict[str, Entry], field: str
+) -> tuple[list[Any], DataType, DataShape]:
+    """Return field values along with their detected type and shape.
+
+    Args:
+        entries: Mapping of uid to entry.
+        field: Field name to extract.
+
+    Returns:
+        Tuple of (values, data type, data shape).
+    """
+    values = get_field_values(entries, field)
+    meta = detect_field_meta(values)
+    return values, meta.data_type, meta.data_shape
+
+
+def _select_best_contextual_field(
+    candidates: list[tuple[str, Any]],
+    *,
+    ctx: EntryContext,
+) -> Any:
+    """Select best contextual field from multiple candidates.
+
+    Priority:
+    1. Fields whose context includes current parameter
+    2. Among those, prefer smaller context size (more specific)
+    3. Fallback: deterministic sort by field name
+    """
+    scored: list[tuple[int, int, str, Any]] = []
+
+    for field_name, value in candidates:
+        _, models, params = parse_contextual_field_name(field_name)
+
+        model_match = models is None or (
+            ctx.model_id is not None and ctx.model_id in models
+        )
+        param_match = params is None or (
+            ctx.parameter_name is not None and ctx.parameter_name in params
+        )
+
+        context_size = (len(models) if models else 0) + (len(params) if params else 0)
+
+        scored.append(
+            (
+                1 if model_match else 0,
+                1 if param_match else 0,
+                -context_size,
+                field_name,
+                value,
+            )
+        )
+
+    scored.sort(key=lambda x: (-x[0], -x[1], -x[2], x[3]))
+
+    return scored[0][4]

diffract/viz/data/filtering.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from .extraction import get_field_value
+from .types import Entry, EntryContext
+
+
+def apply_filter(
+    entries: dict[str, Entry],
+    value_filter: dict[str, tuple[str, Any]],
+) -> dict[str, Entry]:
+    """Apply value-based filtering to entries."""
+    return {
+        uid: entry
+        for uid, entry in entries.items()
+        if _entry_passes_filter(entry, value_filter)
+    }
+
+
+def _entry_passes_filter(
+    entry: Entry,
+    value_filter: dict[str, tuple[str, Any]],
+) -> bool:
+    """Check if an entry passes all filter conditions."""
+    ctx = EntryContext.from_entry(entry)
+    for field_name, (op, threshold) in value_filter.items():
+        value = get_field_value(ctx, field_name)
+        if not _check_condition(value, op, threshold):
+            return False
+    return True
+
+
+def _check_condition(value: Any, op: str, threshold: Any) -> bool:
+    """Check a single filter condition."""
+    if value is None:
+        return False
+
+    if isinstance(value, np.ndarray):
+        value = float(np.nanmean(value))
+
+    ops = {
+        ">": lambda a, b: a > b,
+        "<": lambda a, b: a < b,
+        ">=": lambda a, b: a >= b,
+        "<=": lambda a, b: a <= b,
+        "==": lambda a, b: a == b,
+        "!=": lambda a, b: a != b,
+    }
+
+    op_func = ops.get(op)
+    if op_func is None:
+        raise ValueError(f"Unknown operator: {op}")
+
+    return op_func(value, threshold)