diffract-core 0.2.1__py3-none-any.whl

diffract/core/export/formatters/base.py

@@ -0,0 +1,68 @@
+"""Base utilities for DataFrame formatters.
+
+Provides shared logic for building scalar and aggregate records.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from diffract.core.export.interface import AggregateData, ResultData
+
+SCALAR_COLUMNS = ["model_id", "parameter_name", "parameter_uid", "parameter_type"]
+AGGREGATE_COLUMNS = [
+    "field",
+    "context_models",
+    "context_params",
+    "value",
+]
+
+
+def build_scalar_records(results: ResultData) -> list[dict[str, Any]]:
+    """Build scalar records from parameter results.
+
+    Args:
+        results: Parameter results dictionary.
+
+    Returns:
+        List of scalar records for DataFrame creation.
+    """
+    scalar_records: list[dict[str, Any]] = []
+
+    for param_uid, param_data in results.items():
+        metadata = param_data["metadata"]
+        field_values = param_data["fields"]
+
+        base_meta = {
+            "parameter_uid": param_uid,
+            "model_id": metadata["model_id"],
+            "parameter_name": metadata["name"],
+            "parameter_type": metadata["parameter_type"],
+        }
+
+        scalar_record: dict[str, Any] = {**base_meta}
+
+        for key, value in metadata.items():
+            if key not in ["model_id", "name", "parameter_type"]:
+                scalar_record[f"meta_{key}"] = value
+
+        scalar_record.update(field_values)
+
+        # Only add record if it has data beyond base metadata
+        if len(scalar_record) > len(base_meta):
+            scalar_records.append(scalar_record)
+
+    return scalar_records
+
+
+def build_aggregate_records(aggregates: AggregateData) -> list[dict[str, Any]]:
+    """Build aggregate records from aggregate results.
+
+    Args:
+        aggregates: Aggregate results list.
+
+    Returns:
+        List of aggregate records for DataFrame creation.
+    """
+    # Aggregates are already in the right format, just return them
+    return list(aggregates)

diffract/core/export/formatters/dict_formatter.py

@@ -0,0 +1,38 @@
+"""Dictionary formatter for result export.
+
+Returns structured results with separate scalars and aggregates dictionaries.
+"""
+
+from __future__ import annotations
+
+from diffract.core.export.interface import (
+    AggregateData,
+    IResultFormatter,
+    ResultData,
+    StructuredExportResult,
+)
+
+
+class DictFormatter(IResultFormatter):
+    """Return results as structured dictionaries."""
+
+    def format_results(
+        self,
+        param_results: ResultData,
+        aggregate_results: AggregateData,
+        _fields: tuple[str, ...],
+    ) -> StructuredExportResult[dict]:
+        """Return results as structured dict with scalars and aggregates.
+
+        Args:
+            param_results: Parameter results dictionary with scalar fields.
+            aggregate_results: Aggregate results list with contextual fields.
+            _fields: Field names (unused for dict format).
+
+        Returns:
+            StructuredExportResult with scalars dict and aggregates list.
+        """
+        return StructuredExportResult(
+            scalars=param_results,
+            aggregates=aggregate_results,
+        )

diffract/core/export/formatters/json_formatter.py

@@ -0,0 +1,58 @@
+"""JSON formatter for result export.
+
+Serializes structured results to a JSON string with a custom serializer that
+converts array-like objects (e.g., numpy arrays, tensors) via tolist()
+when available, and falls back to string representation otherwise.
+"""
+
+from __future__ import annotations
+
+import json
+
+from diffract.core.export.interface import (
+    AggregateData,
+    IResultFormatter,
+    ResultData,
+)
+
+
+class JsonFormatter(IResultFormatter):
+    """Convert results to a pretty-printed JSON string."""
+
+    def format_results(
+        self,
+        param_results: ResultData,
+        aggregate_results: AggregateData,
+        _fields: tuple[str, ...],
+    ) -> str:
+        """Convert results to formatted JSON string.
+
+        Args:
+            param_results: Parameter results dictionary with scalar fields.
+            aggregate_results: Aggregate results list with contextual fields.
+            _fields: Field names (unused for JSON format).
+
+        Returns:
+            JSON string representation with scalars and aggregates.
+        """
+        combined = {
+            "scalars": param_results,
+            "aggregates": aggregate_results,
+        }
+        return json.dumps(combined, indent=2, default=self._json_serializer)
+
+    def _json_serializer(self, obj: object) -> object:
+        """Custom JSON serializer for non-native types.
+
+        Uses tolist() for array-like structures to preserve value structure.
+        Falls back to str() for other objects.
+
+        Args:
+            obj: Object to serialize.
+
+        Returns:
+            JSON-serializable representation of the object.
+        """
+        if hasattr(obj, "tolist"):
+            return obj.tolist()
+        return str(obj)

diffract/core/export/formatters/list_formatter.py

@@ -0,0 +1,41 @@
+"""List formatter for result export.
+
+Returns structured results with scalar records as list-of-dicts and
+aggregates as list-of-dicts.
+"""
+
+from __future__ import annotations
+
+from diffract.core.export.interface import (
+    AggregateData,
+    IResultFormatter,
+    ResultData,
+    StructuredExportResult,
+)
+
+from .base import build_aggregate_records, build_scalar_records
+
+
+class ListFormatter(IResultFormatter):
+    """Return results as list-based records."""
+
+    def format_results(
+        self,
+        param_results: ResultData,
+        aggregate_results: AggregateData,
+        _fields: tuple[str, ...],
+    ) -> StructuredExportResult[list[dict]]:
+        """Return results as lists of records.
+
+        Args:
+            param_results: Parameter results dictionary with scalar fields.
+            aggregate_results: Aggregate results list with contextual fields.
+            _fields: Field names (unused for list format).
+
+        Returns:
+            StructuredExportResult with scalars list and aggregates list.
+        """
+        return StructuredExportResult(
+            scalars=build_scalar_records(param_results),
+            aggregates=build_aggregate_records(aggregate_results),
+        )

diffract/core/export/formatters/pandas_formatter.py

@@ -0,0 +1,86 @@
+"""Pandas DataFrame formatter for result export.
+
+Creates structured DataFrames separating scalar per-parameter fields from
+contextual/aggregation fields. If pandas is unavailable, raises ImportError.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Self
+
+import diffract.core.utils.imports as import_utils
+from diffract.core.export.interface import (
+    AggregateData,
+    IResultFormatter,
+    ResultData,
+    StructuredExportResult,
+)
+
+from .base import (
+    AGGREGATE_COLUMNS,
+    SCALAR_COLUMNS,
+    build_aggregate_records,
+    build_scalar_records,
+)
+
+logger = logging.getLogger(__name__)
+
+
+if not import_utils.is_available("pandas"):
+    logger.debug("Failed to import pandas, disabling corresponding formatters")
+
+    class PandasFormatter(IResultFormatter):
+        """Stub formatter that raises when pandas is unavailable."""
+
+        def __new__(cls, *_args: Any, **_kwargs: Any) -> Self:
+            """Raise ImportError because pandas is not installed."""
+            msg = "pandas package not available"
+            raise ImportError(msg)
+
+else:
+    pd = import_utils.require("pandas")
+
+    class PandasFormatter(IResultFormatter):
+        """Convert results to structured DataFrames separating scalars/aggregates."""
+
+        def format_results(
+            self,
+            param_results: ResultData,
+            aggregate_results: AggregateData,
+            fields: tuple[str, ...],
+        ) -> StructuredExportResult[pd.DataFrame]:
+            """Convert results to structured pandas DataFrames.
+
+            Args:
+                param_results: Parameter results dictionary with scalar fields.
+                aggregate_results: Aggregate results list with contextual fields.
+                fields: Field names requested for export.
+
+            Returns:
+                StructuredExportResult with scalars and aggregates DataFrames.
+            """
+            if not param_results and not aggregate_results:
+                return StructuredExportResult(
+                    scalars=pd.DataFrame(columns=[*SCALAR_COLUMNS, *fields]),
+                    aggregates=pd.DataFrame(columns=AGGREGATE_COLUMNS),
+                )
+
+            scalar_records = build_scalar_records(param_results)
+            aggregate_records = build_aggregate_records(aggregate_results)
+
+            scalars_df = (
+                pd.DataFrame(scalar_records)
+                if scalar_records
+                else pd.DataFrame(columns=[*SCALAR_COLUMNS, *fields])
+            )
+            aggregates_df = (
+                pd.DataFrame(aggregate_records)
+                if aggregate_records
+                else pd.DataFrame(columns=AGGREGATE_COLUMNS)
+            )
+
+            return StructuredExportResult(
+                scalars=scalars_df,
+                aggregates=aggregates_df,
+            )

diffract/core/export/formatters/polars_formatter.py

@@ -0,0 +1,86 @@
+"""Polars DataFrame formatter for result export.
+
+Creates structured DataFrames separating scalar per-parameter fields from
+contextual/aggregation fields. If polars is unavailable, raises ImportError.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Self
+
+import diffract.core.utils.imports as import_utils
+from diffract.core.export.interface import (
+    AggregateData,
+    IResultFormatter,
+    ResultData,
+    StructuredExportResult,
+)
+
+from .base import (
+    AGGREGATE_COLUMNS,
+    SCALAR_COLUMNS,
+    build_aggregate_records,
+    build_scalar_records,
+)
+
+logger = logging.getLogger(__name__)
+
+
+if not import_utils.is_available("polars"):
+    logger.debug("Failed to import polars, disabling corresponding formatters")
+
+    class PolarsFormatter(IResultFormatter):
+        """Stub formatter that raises when polars is unavailable."""
+
+        def __new__(cls, *_args: Any, **_kwargs: Any) -> Self:
+            """Raise ImportError because polars is not installed."""
+            msg = "polars package not available"
+            raise ImportError(msg)
+
+else:
+    pl = import_utils.require("polars")
+
+    class PolarsFormatter(IResultFormatter):
+        """Convert results to structured DataFrames separating scalars/aggregates."""
+
+        def format_results(
+            self,
+            param_results: ResultData,
+            aggregate_results: AggregateData,
+            fields: tuple[str, ...],
+        ) -> StructuredExportResult[pl.DataFrame]:
+            """Convert results to structured polars DataFrames.
+
+            Args:
+                param_results: Parameter results dictionary with scalar fields.
+                aggregate_results: Aggregate results list with contextual fields.
+                fields: Field names requested for export.
+
+            Returns:
+                StructuredExportResult with scalars and aggregates DataFrames.
+            """
+            if not param_results and not aggregate_results:
+                return StructuredExportResult(
+                    scalars=pl.DataFrame({c: [] for c in [*SCALAR_COLUMNS, *fields]}),
+                    aggregates=pl.DataFrame({c: [] for c in AGGREGATE_COLUMNS}),
+                )
+
+            scalar_records = build_scalar_records(param_results)
+            aggregate_records = build_aggregate_records(aggregate_results)
+
+            scalars_df = (
+                pl.DataFrame(scalar_records)
+                if scalar_records
+                else pl.DataFrame({c: [] for c in [*SCALAR_COLUMNS, *fields]})
+            )
+            aggregates_df = (
+                pl.DataFrame(aggregate_records)
+                if aggregate_records
+                else pl.DataFrame({c: [] for c in AGGREGATE_COLUMNS})
+            )
+
+            return StructuredExportResult(
+                scalars=scalars_df,
+                aggregates=aggregates_df,
+            )

diffract/core/export/formatters/registry.py

@@ -0,0 +1,84 @@
+"""Formatter registry and factory for result export.
+
+Maintains a mapping from format names to formatter instances and provides
+helpers to register custom formatters and retrieve formatters by name.
+"""
+
+from __future__ import annotations
+
+import difflib
+from typing import TYPE_CHECKING
+
+import diffract.core.utils.imports as import_utils
+
+from .dict_formatter import DictFormatter
+from .json_formatter import JsonFormatter
+from .list_formatter import ListFormatter
+
+if TYPE_CHECKING:
+    from diffract.core.export.interface import IResultFormatter
+
+# Formats backed by optional dependencies, mapped to the extra providing them
+_OPTIONAL_FORMAT_EXTRAS: dict[str, str] = {
+    "pandas": "pandas",
+    "polars": "polars",
+}
+
+# Built-in formatter instances
+FORMATTERS: dict[str, IResultFormatter] = {
+    "dict": DictFormatter(),
+    "json": JsonFormatter(),
+    "list": ListFormatter(),
+}
+
+if import_utils.is_available("pandas"):
+    from .pandas_formatter import PandasFormatter
+
+    FORMATTERS["pandas"] = PandasFormatter()  # type: ignore[arg-type]
+
+if import_utils.is_available("polars"):
+    from .polars_formatter import PolarsFormatter
+
+    FORMATTERS["polars"] = PolarsFormatter()  # type: ignore[arg-type]
+
+
+def register_formatter(name: str, formatter: IResultFormatter) -> None:
+    """Register a custom formatter under a given name.
+
+    Overwrites an existing formatter if the name is already used.
+
+    Args:
+        name: Unique format name.
+        formatter: Formatter instance to register.
+    """
+    FORMATTERS[name] = formatter
+
+
+def get_formatter(name: str) -> IResultFormatter:
+    """Return a formatter by name or raise ValueError if unknown.
+
+    Args:
+        name: Format name to lookup.
+
+    Returns:
+        Formatter instance associated with the name.
+
+    Raises:
+        ValueError: If name is not registered.
+    """
+    try:
+        return FORMATTERS[name]
+    except KeyError as e:
+        if name in _OPTIONAL_FORMAT_EXTRAS:
+            extra = _OPTIONAL_FORMAT_EXTRAS[name]
+            msg = (
+                f"Export format '{name}' requires the optional dependency "
+                f"'{extra}'. Install it with: "
+                f'pip install "diffract-core[{extra}]".'
+            )
+        else:
+            known = ", ".join(sorted(FORMATTERS))
+            close = difflib.get_close_matches(name, FORMATTERS, n=3, cutoff=0.6)
+            hint = f" Did you mean: {', '.join(close)}?" if close else ""
+            msg = f"Unsupported format '{name}'. Known: {known}.{hint}"
+        raise ValueError(msg) from e

diffract/core/export/interface.py

@@ -0,0 +1,105 @@
+"""Interfaces (protocols) for result export and formatting.
+
+Defines:
+    - IResultFormatter: Converts structured results to a target format.
+    - IResultExporter: Extracts requested fields from parameters and aggregates,
+      then delegates formatting to an IResultFormatter.
+
+ResultData structure (parameters):
+    {
+        "<param_uid>": {
+            "metadata": { ... },
+            "fields": { "field": value, ... }
+        },
+        ...
+    }
+
+AggregateData structure:
+    [
+        {
+            "field": "field_name",
+            "context_models": ("model1", "model2"),
+            "context_params": ("param1",),
+            "value": ...,
+        },
+        ...
+    ]
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar, runtime_checkable
+
+if TYPE_CHECKING:
+    from diffract.core.data.nn.aggregates.view import AggregateView
+    from diffract.core.data.nn.params.interface import IParameterView
+
+ResultData = dict[str, dict[str, Any]]
+AggregateData = list[dict[str, Any]]
+
+T = TypeVar("T")
+
+
+@dataclass
+class StructuredExportResult[T]:
+    """Container for structured export with separate scalar and aggregate data.
+
+    Scalars are per-parameter fields (e.g., frob_norm, stable_rank).
+    Aggregates are cross-entity fields from aggregation kernels (e.g., l_overlap).
+
+    Attributes:
+        scalars: DataFrame/dict with per-parameter metrics indexed by parameter.
+        aggregates: DataFrame/dict with aggregation results including context info.
+    """
+
+    scalars: T
+    aggregates: T
+
+
+@runtime_checkable
+class IResultFormatter(Protocol):
+    """Protocol for converting raw results to specific formats."""
+
+    def format_results(
+        self,
+        param_results: ResultData,
+        aggregate_results: AggregateData,
+        fields: tuple[str, ...],
+    ) -> Any:
+        """Convert raw results to the target format.
+
+        Args:
+            param_results: Parameter results dictionary with scalar fields.
+            aggregate_results: Aggregate results list with contextual fields.
+            fields: Field names requested for export.
+
+        Returns:
+            Formatted results in the target format (StructuredExportResult, dict, etc.).
+        """
+        ...
+
+
+@runtime_checkable
+class IResultExporter(Protocol):
+    """Protocol for exporting results from parameter and aggregate repositories."""
+
+    def export_results(
+        self,
+        *fields: tuple[str, ...],
+        parameters: IParameterView,
+        aggregates: AggregateView | None,
+        formatter: IResultFormatter,
+    ) -> Any:
+        """Export results from parameters and aggregates using the specified formatter.
+
+        Args:
+            *fields: Field names to export.
+            parameters: Parameter collection to export scalar fields from.
+            aggregates: Aggregate view to export contextual fields from.
+            formatter: Formatter to use for output conversion.
+
+        Returns:
+            Formatted results from the formatter.
+        """
+        ...

diffract/core/parallel/__init__.py

@@ -0,0 +1,21 @@
+"""Parallel execution utilities for compute subsystem."""
+
+from .containers import ParallelSingletonContainer
+from .execution import map_maybe_parallel
+from .runtime import (
+    ParallelCalibration,
+    ParallelContext,
+    calibrate_thread_pool_overhead,
+    get_thread_pool_calibration,
+    should_parallelize,
+)
+
+__all__ = [
+    "ParallelCalibration",
+    "ParallelContext",
+    "ParallelSingletonContainer",
+    "calibrate_thread_pool_overhead",
+    "get_thread_pool_calibration",
+    "map_maybe_parallel",
+    "should_parallelize",
+]

diffract/core/parallel/containers.py

@@ -0,0 +1,57 @@
+"""Dependency injection container for pool-per-method executors."""
+
+from __future__ import annotations
+
+import concurrent.futures
+import os
+from typing import Any
+
+from dependency_injector import containers, providers
+
+from .runtime import ParallelContext, get_thread_pool_calibration
+
+
+def _normalize_workers(max_workers: int | None) -> int:
+    if max_workers is None:
+        return max(1, os.cpu_count() or 1)
+    return max(1, int(max_workers))
+
+
+def _thread_pool_context_resource(*, workers: int) -> Any:
+    executor = concurrent.futures.ThreadPoolExecutor(max_workers=workers)
+    try:
+        calibration = get_thread_pool_calibration(workers)
+        yield ParallelContext(
+            executor=executor, calibration=calibration, workers=workers
+        )
+    finally:
+        executor.shutdown(wait=True, cancel_futures=True)
+
+
+def _process_pool_context_resource(*, workers: int) -> Any:
+    executor = concurrent.futures.ProcessPoolExecutor(max_workers=workers)
+    try:
+        yield executor
+    finally:
+        executor.shutdown(wait=True, cancel_futures=True)
+
+
+class ParallelSingletonContainer(containers.DeclarativeContainer):
+    """Singleton container for project-wide parallelism resources."""
+
+    config = providers.Configuration()
+
+    max_workers = providers.Callable(_normalize_workers, config.thread_pool.max_workers)
+    process_max_workers = providers.Callable(
+        _normalize_workers, config.process_pool.max_workers
+    )
+
+    thread_pool_context = providers.Resource(
+        _thread_pool_context_resource,
+        workers=max_workers,
+    )
+
+    process_pool_context = providers.Resource(
+        _process_pool_context_resource,
+        workers=process_max_workers,
+    )