diffract-core 0.2.1__py3-none-any.whl

diffract/core/data/nn/aggregates/metadata.py

@@ -0,0 +1,96 @@
+"""Metadata for aggregate entities.
+
+This module provides the AggregateMetadata class for representing
+aggregated data between parameters/models computed by aggregation kernels.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any, Self
+
+from diffract.core.utils.hashing import get_unique_id
+
+
+@dataclass(frozen=True, kw_only=True)
+class AggregateMetadata:
+    """Immutable metadata container for aggregate entities.
+
+    An aggregate represents computed data involving multiple models or parameters,
+    such as overlap metrics between two models or correlation between parameters.
+
+    Attributes:
+        uid: Unique identifier automatically generated if not provided.
+        field_name: Base field name of the computed aggregate (e.g., "l_overlap").
+        context_models: Tuple of model IDs participating in this aggregate.
+        context_params: Tuple of parameter names participating in this aggregate.
+    """
+
+    uid: str = field(default_factory=get_unique_id)
+    field_name: str
+    context_models: tuple[str, ...] = ()
+    context_params: tuple[str, ...] = ()
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize metadata to dictionary.
+
+        Returns:
+            Dictionary representation suitable for storage/reconstruction.
+        """
+        return {
+            "uid": self.uid,
+            "field_name": self.field_name,
+            "context_models": json.dumps(list(self.context_models)),
+            "context_params": json.dumps(list(self.context_params)),
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Self:
+        """Deserialize metadata from dictionary.
+
+        Args:
+            data: Dictionary with metadata fields.
+
+        Returns:
+            AggregateMetadata instance.
+        """
+        # Handle both list format (from to_dict) and JSON string format (from index)
+        context_models = data.get("context_models", [])
+        if isinstance(context_models, str):
+            context_models = json.loads(context_models)
+
+        context_params = data.get("context_params", [])
+        if isinstance(context_params, str):
+            context_params = json.loads(context_params)
+
+        return cls(
+            uid=data["uid"],
+            field_name=data["field_name"],
+            context_models=tuple(context_models),
+            context_params=tuple(context_params),
+        )
+
+    @classmethod
+    def create_uid_from_context(
+        cls,
+        field_name: str,
+        context_models: tuple[str, ...],
+        context_params: tuple[str, ...],
+    ) -> str:
+        """Create a deterministic UID from the aggregate context.
+
+        This allows deduplication of aggregates with the same context.
+        Uses the same format as contextual field names in constants.py.
+
+        Args:
+            field_name: Base field name of the aggregate.
+            context_models: Tuple of model IDs.
+            context_params: Tuple of parameter names.
+
+        Returns:
+            A deterministic UID string based on the context.
+        """
+        from diffract.core.constants import format_contextual_field_name
+
+        return format_contextual_field_name(field_name, context_models, context_params)

diffract/core/data/nn/aggregates/proxy.py

@@ -0,0 +1,57 @@
+"""Proxy for aggregate entities.
+
+This module provides the AggregateProxy class for accessing
+aggregated data stored in the aggregates table.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from diffract.core.constants import TABLE_AGGREGATES
+from diffract.core.data.proxy import DataProxy
+
+from .metadata import AggregateMetadata
+
+if TYPE_CHECKING:
+    from .repository import AggregateRepository
+
+
+@dataclass(kw_only=True)
+class AggregateProxy(DataProxy[AggregateMetadata]):
+    """Lazy-loading proxy for aggregate data.
+
+    Extends the generic DataProxy with aggregate-specific configuration.
+    Aggregates are stored in the TABLE_AGGREGATES table and represent
+    computed aggregations between models/parameters.
+
+    Attributes:
+        meta: Immutable aggregate metadata.
+        _repository: Repository that owns storage/cache managers.
+    """
+
+    meta: AggregateMetadata
+    _repository: AggregateRepository = field(repr=False)
+
+    @classmethod
+    def get_table(cls) -> str:
+        """Return the storage table name for aggregates."""
+        return TABLE_AGGREGATES
+
+    @classmethod
+    def create_and_store(
+        cls,
+        meta: AggregateMetadata,
+        repository: AggregateRepository,
+    ) -> AggregateProxy:
+        """Create aggregate proxy and store metadata in index.
+
+        Args:
+            meta: Aggregate metadata.
+            repository: Repository owning storage, cache, and metadata managers.
+
+        Returns:
+            New aggregate proxy with metadata stored.
+        """
+        return super().create_and_store(meta=meta, repository=repository)

diffract/core/data/nn/aggregates/repository.py

@@ -0,0 +1,87 @@
+"""Repository for aggregate entities.
+
+This module provides the AggregateRepository class for managing
+aggregate data storage and membership.
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+from diffract.core.constants import TABLE_AGGREGATES
+from diffract.core.data.repository import DataRepository
+
+from .metadata import AggregateMetadata
+from .proxy import AggregateProxy
+from .view import AggregateView
+
+
+class AggregateRepository(DataRepository[AggregateMetadata, AggregateProxy]):
+    """Repository owning storage/cache/metadata and managing aggregate membership.
+
+    Extends the generic DataRepository with aggregate-specific configuration
+    for metadata class, proxy class, view class, storage table, and metadata schema.
+    """
+
+    METADATA_CLASS = AggregateMetadata
+    PROXY_CLASS = AggregateProxy
+    VIEW_CLASS = AggregateView
+    TABLE = TABLE_AGGREGATES
+
+    # Schema for MetadataIndex
+    METADATA_COLUMNS: ClassVar[dict[str, type]] = {
+        "field_name": str,
+        "context_models": str,  # JSON-serialized tuple
+        "context_params": str,  # JSON-serialized tuple
+    }
+    METADATA_INDEXES: ClassVar[list[str]] = ["field_name"]
+
+    def get_or_create(
+        self,
+        field_name: str,
+        context_models: tuple[str, ...],
+        context_params: tuple[str, ...],
+    ) -> AggregateProxy:
+        """Get existing aggregate by context or create a new one.
+
+        This method provides deduplication for aggregates with the same context.
+
+        Args:
+            field_name: Base field name of the aggregate.
+            context_models: Tuple of model IDs participating in this aggregate.
+            context_params: Tuple of parameter names participating in this aggregate.
+
+        Returns:
+            Existing or newly created AggregateProxy.
+        """
+        # Create deterministic UID
+        uid = AggregateMetadata.create_uid_from_context(
+            field_name=field_name,
+            context_models=context_models,
+            context_params=context_params,
+        )
+
+        # Check if already exists in cache or index
+        if uid in self._proxy_cache:
+            return self._proxy_cache[uid]
+
+        if self.has_uid(uid):
+            return self.get_proxy(uid)
+
+        # Create new aggregate
+        meta = AggregateMetadata(
+            uid=uid,
+            field_name=field_name,
+            context_models=context_models,
+            context_params=context_params,
+        )
+
+        return AggregateProxy.create_and_store(meta=meta, repository=self)
+
+    def create_view(self) -> AggregateView:
+        """Create a view over all aggregates in this repository.
+
+        Returns:
+            AggregateView containing all aggregates.
+        """
+        return AggregateView(repository=self, uids=None)

diffract/core/data/nn/aggregates/view.py

@@ -0,0 +1,307 @@
+"""View for aggregate entities.
+
+This module provides the AggregateView class for batch operations
+on filtered aggregate collections.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Iterator
+from typing import TYPE_CHECKING
+
+from diffract.core.constants import REGEX_PREFIX, TABLE_AGGREGATES
+from diffract.core.data.interface import EntityUID
+from diffract.core.data.utils import build_matcher
+from diffract.core.data.view import DataView
+
+from .metadata import AggregateMetadata
+from .proxy import AggregateProxy
+
+if TYPE_CHECKING:
+    from .repository import AggregateRepository
+
+
+class AggregateView(DataView[AggregateMetadata, AggregateProxy]):
+    """A numpy-like view over a subset of aggregates owned by a repository.
+
+    Extends the generic DataView with aggregate-specific filtering methods.
+    Uses SQL-based filtering via MetadataIndex for efficiency.
+    """
+
+    def __init__(
+        self,
+        *,
+        repository: AggregateRepository,
+        uids: list[EntityUID] | None = None,
+    ) -> None:
+        super().__init__(repository=repository, uids=uids)
+
+    def __iter__(self) -> Iterator[AggregateProxy]:
+        """Iterate aggregate proxies in view order."""
+        return super().__iter__()
+
+    def _sort_in_place(self) -> None:
+        """Sort UIDs by field_name, context for deterministic ordering."""
+        uids = self._ensure_uids()
+        uids.sort(
+            key=lambda uid: (
+                self._repository.get_proxy(uid).meta.field_name,
+                ",".join(sorted(self._repository.get_proxy(uid).meta.context_models)),
+                ",".join(sorted(self._repository.get_proxy(uid).meta.context_params)),
+                uid,
+            )
+        )
+        self._sorted = True
+
+    def filter_by_field_name(
+        self, *names: str, inverse_mask: bool = False
+    ) -> AggregateView:
+        """Filter aggregates by field name.
+
+        Args:
+            *names: Field names to match. Use "re:" prefix for regex patterns.
+            inverse_mask: If True, return aggregates NOT matching the names.
+
+        Returns:
+            New view with filtered aggregates.
+        """
+        # Separate exact matches from regex patterns
+        exact_names: list[str] = []
+        regex_patterns: list[re.Pattern[str]] = []
+        for name in names:
+            if name.startswith(REGEX_PREFIX):
+                regex_patterns.append(re.compile(name.removeprefix(REGEX_PREFIX)))
+            else:
+                exact_names.append(name)
+
+        # If only exact matches and no inverse, use SQL
+        if exact_names and not regex_patterns and not inverse_mask:
+            current_uids = self._uids
+            if current_uids is None:
+                # Query directly from index
+                filtered = self._repository.metadata_index.query(
+                    TABLE_AGGREGATES,
+                    where_in={"field_name": exact_names},
+                )
+            else:
+                # Filter within current UIDs
+                filtered = self._repository.metadata_index.query(
+                    TABLE_AGGREGATES,
+                    where_in={"field_name": exact_names, "uid": current_uids},
+                )
+            return AggregateView(repository=self._repository, uids=filtered)
+
+        # Fall back to in-memory filtering for regex or inverse
+        if not self._sorted:
+            self._sort_in_place()
+
+        uids = self._ensure_uids()
+        matches = build_matcher(names)
+
+        if not inverse_mask:
+            filtered = [
+                uid
+                for uid in uids
+                if matches(self._repository.get_proxy(uid).meta.field_name)
+            ]
+        else:
+            filtered = [
+                uid
+                for uid in uids
+                if not matches(self._repository.get_proxy(uid).meta.field_name)
+            ]
+
+        return AggregateView(repository=self._repository, uids=filtered)
+
+    def filter_by_context_models(
+        self, *model_ids: str, require_all: bool = False
+    ) -> AggregateView:
+        """Filter aggregates by participating models.
+
+        Args:
+            *model_ids: Model IDs that must be in the aggregate's context_models.
+            require_all: If True, all model_ids must be present. If False,
+                         at least one must be present.
+
+        Returns:
+            New view with filtered aggregates.
+        """
+        exact_names: set[str] = set()
+        regex_patterns: list[re.Pattern[str]] = []
+        for model_id in model_ids:
+            if model_id.startswith(REGEX_PREFIX):
+                regex_patterns.append(re.compile(model_id.removeprefix(REGEX_PREFIX)))
+            else:
+                exact_names.add(model_id)
+
+        if not self._sorted:
+            self._sort_in_place()
+
+        uids = self._ensure_uids()
+
+        def _matches(context_models: tuple[str, ...]) -> bool:
+            context_set = set(context_models)
+
+            if exact_names:
+                if require_all:
+                    if not exact_names <= context_set:
+                        return False
+                elif not (exact_names & context_set):
+                    return False
+
+            if regex_patterns:
+                if require_all:
+                    for pattern in regex_patterns:
+                        if not any(
+                            pattern.fullmatch(cm) is not None for cm in context_models
+                        ):
+                            return False
+                elif not any(
+                    any(pattern.fullmatch(cm) is not None for pattern in regex_patterns)
+                    for cm in context_models
+                ):
+                    return False
+
+            return True
+
+        filtered = [
+            uid
+            for uid in uids
+            if _matches(self._repository.get_proxy(uid).meta.context_models)
+        ]
+
+        return AggregateView(repository=self._repository, uids=filtered)
+
+    def filter_by_context_params(
+        self, *param_names: str, require_all: bool = False
+    ) -> AggregateView:
+        """Filter aggregates by participating parameters.
+
+        Args:
+            *param_names: Parameter names that must be in the aggregate's
+                context_params.
+            require_all: If True, all param_names must be present. If False,
+                         at least one must be present.
+
+        Returns:
+            New view with filtered aggregates.
+        """
+        exact_names: set[str] = set()
+        regex_patterns: list[re.Pattern[str]] = []
+        for param_name in param_names:
+            if param_name.startswith(REGEX_PREFIX):
+                regex_patterns.append(re.compile(param_name.removeprefix(REGEX_PREFIX)))
+            else:
+                exact_names.add(param_name)
+
+        if not self._sorted:
+            self._sort_in_place()
+
+        uids = self._ensure_uids()
+
+        def _matches(context_params: tuple[str, ...]) -> bool:
+            context_set = set(context_params)
+
+            if exact_names:
+                if require_all:
+                    if not exact_names <= context_set:
+                        return False
+                elif not (exact_names & context_set):
+                    return False
+
+            if regex_patterns:
+                if require_all:
+                    for pattern in regex_patterns:
+                        if not any(
+                            pattern.fullmatch(cp) is not None for cp in context_params
+                        ):
+                            return False
+                elif not any(
+                    any(pattern.fullmatch(cp) is not None for pattern in regex_patterns)
+                    for cp in context_params
+                ):
+                    return False
+
+            return True
+
+        filtered = [
+            uid
+            for uid in uids
+            if _matches(self._repository.get_proxy(uid).meta.context_params)
+        ]
+
+        return AggregateView(repository=self._repository, uids=filtered)
+
+    def filter_by_exact_context(
+        self,
+        *,
+        models: tuple[str, ...] | None = None,
+        params: tuple[str, ...] | None = None,
+    ) -> AggregateView:
+        """Filter aggregates by exact context match.
+
+        Unlike filter_by_context_models/params which use set intersection,
+        this method requires exact equality of the context tuples.
+
+        Args:
+            models: Exact context_models tuple to match. If None, any models match.
+            params: Exact context_params tuple to match. If None, any params match.
+
+        Returns:
+            New view with aggregates matching the exact context.
+        """
+        if not self._sorted:
+            self._sort_in_place()
+
+        uids = self._ensure_uids()
+
+        def _matches(proxy: AggregateProxy) -> bool:
+            if models is not None and proxy.meta.context_models != models:
+                return False
+            return not (params is not None and proxy.meta.context_params != params)
+
+        filtered = [uid for uid in uids if _matches(self._repository.get_proxy(uid))]
+
+        return AggregateView(repository=self._repository, uids=filtered)
+
+    def erase_fields(self, *fields: str) -> None:
+        """Erase fields from all aggregates in this view."""
+        uids = self._ensure_uids()
+        if not uids:
+            return
+
+        with self:
+            for field in fields:
+                self._repository.storage_manager.erase_field_for_all(
+                    field, table=TABLE_AGGREGATES
+                )
+                if self._repository.cache_manager is not None:
+                    self._repository.cache_manager.erase_field_for_all(field)
+
+    def clear(self, erase: bool = False) -> None:
+        """Clear this view and optionally erase corresponding data.
+
+        Args:
+            erase: If True, also erase underlying storage data.
+                   If False, only clear membership (metadata index).
+        """
+        uids = self._ensure_uids()
+
+        if erase and self._repository.cache_manager is not None:
+            self._repository.cache_manager.clear()
+
+        with self._repository:
+            for uid in list(uids):
+                self._repository._proxy_cache.pop(uid, None)
+                self._repository.metadata_index.delete(TABLE_AGGREGATES, uid)
+                if erase:
+                    self._repository.storage_manager.erase_obj(
+                        uid, table=TABLE_AGGREGATES
+                    )
+
+        if self._uids is not None:
+            self._uids.clear()
+        else:
+            self._uids = []
+        self._sorted = True

diffract/core/data/nn/containers.py

@@ -0,0 +1,86 @@
+"""Dependency injection containers for model parameter management.
+
+This module provides dependency injection configuration for model parameter
+extraction and collection management using the dependency-injector library.
+It integrates parameter extractors, storage managers, metadata index, and
+cache managers for comprehensive parameter lifecycle management.
+
+The container provides factory methods for creating parameter extractors
+and initializing parameter collections, ensuring proper dependency
+injection and configuration management.
+
+Example:
+    >>> container = ModelParametersContainer()
+    >>> container.config.extractor.framework.from_value("pytorch")
+    >>> container.storage_manager.override(storage_instance)
+    >>> container.metadata_index.override(metadata_instance)
+    >>> container.cache_manager.override(cache_instance)
+    >>> collection = container.parameter_repository()
+"""
+
+from __future__ import annotations
+
+from dependency_injector import containers, providers
+
+from diffract.core.data.metadata.interface import IMetadataIndex
+from diffract.core.parallel import ParallelContext
+from diffract.core.storage.interface import IStorageManager
+from diffract.core.utils.build import build_with_defaults
+
+from .aggregates.repository import AggregateRepository
+from .extractors.factory import create_extractor
+from .params.repository import ParameterRepository
+
+
+class ModelParametersContainer(containers.DeclarativeContainer):
+    """Dependency injection container for model parameter components.
+
+    Provides comprehensive dependency injection setup for parameter extraction,
+    collection management, and data lifecycle operations. The container manages
+    the integration between parameter extractors, storage systems, metadata index,
+    and caching mechanisms.
+
+    The container uses external dependencies for storage, metadata index, and
+    cache managers, allowing flexible backend configuration while maintaining
+    consistent parameter extraction and collection interfaces.
+
+    Attributes:
+        config: Configuration provider for parameter extraction settings.
+        storage_manager: External dependency for persistent storage.
+        metadata_index: External dependency for structured metadata.
+        cache_manager: External dependency for temporary caching.
+        extractor_factory: Factory for creating parameter extractors.
+        parameter_repository: Singleton for parameter repository.
+        aggregate_repository: Singleton for aggregate repository.
+    """
+
+    config = providers.Configuration()
+
+    # External dependencies (injected from parent containers)
+    storage_manager = providers.Dependency(instance_of=IStorageManager)
+    metadata_index = providers.Dependency(instance_of=IMetadataIndex)
+    cache_manager = providers.Dependency()
+    parallel = providers.Dependency(instance_of=ParallelContext)
+
+    parameter_repository = providers.Singleton(
+        ParameterRepository.initialize,
+        storage_manager=storage_manager,
+        metadata_index=metadata_index,
+        cache_manager=cache_manager,
+        parallel=parallel,
+    )
+
+    aggregate_repository = providers.Singleton(
+        AggregateRepository.initialize,
+        storage_manager=storage_manager,
+        metadata_index=metadata_index,
+        cache_manager=cache_manager,
+        parallel=parallel,
+    )
+
+    # Parameter extractor factories
+    extractor_factory = providers.Factory(
+        build_with_defaults,
+        create_extractor,
+        config.extractor.as_(dict),
+    )

diffract/core/data/nn/extractors/__init__.py

@@ -0,0 +1,49 @@
+"""Parameter extraction framework for neural network models.
+
+This module provides a comprehensive framework for extracting parameters from
+neural network models across different deep learning frameworks. It uses a
+handler-based architecture for extensible and framework-specific parameter
+processing.
+
+Key Features:
+    - Framework-agnostic parameter extraction interface
+    - Handler-based processing for different parameter types
+    - Factory pattern for automatic extractor creation
+    - Override system for custom parameter metadata
+    - Extensible architecture for new frameworks and parameter types
+
+Supported Frameworks:
+    - NumPy: dict[str, numpy.ndarray] mappings (no framework required)
+    - PyTorch: nn.Module and state_dict extraction
+    - Extensible design for additional frameworks
+
+Example:
+    >>> from diffract.core.data.nn.extractors import create_extractor
+    >>> extractor = create_extractor(pytorch_model)
+    >>> parameters = extractor.extract_parameters(storage, cache)
+"""
+
+from .base import BaseParameterExtractor, ExtractorOverrides, ParameterOverrides
+from .factory import create_extractor, get_supported_frameworks, get_supported_types
+from .flax import FlaxParamsExtractor
+from .interface import IParameterExtractor
+from .numpy import NumpyDictExtractor
+from .onnx import OnnxModelExtractor
+from .tensorflow import TensorFlowModelExtractor
+from .torch import TorchModuleExtractor, TorchStateDictExtractor
+
+__all__ = [
+    "BaseParameterExtractor",
+    "ExtractorOverrides",
+    "FlaxParamsExtractor",
+    "IParameterExtractor",
+    "NumpyDictExtractor",
+    "OnnxModelExtractor",
+    "ParameterOverrides",
+    "TensorFlowModelExtractor",
+    "TorchModuleExtractor",
+    "TorchStateDictExtractor",
+    "create_extractor",
+    "get_supported_frameworks",
+    "get_supported_types",
+]