diffbio 0.1.0__py3-none-any.whl

diffbio/sources/indexed_embeddings.py

@@ -0,0 +1,128 @@
+"""Indexed embedding-artifact sources with strict row-identity alignment."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, cast
+
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+
+from diffbio.sources.embeddings import EmbeddingArtifactSource, EmbeddingArtifactSourceConfig
+
+
+@dataclass(frozen=True)
+class IndexedEmbeddingSourceConfig(EmbeddingArtifactSourceConfig):
+    """Configuration for embedding sources that persist row identities."""
+
+    row_id_key: str | None = None
+
+    def __post_init__(self) -> None:
+        """Validate the required row-identity field name."""
+        super().__post_init__()
+        if self.row_id_key is None:
+            raise ValueError("row_id_key is required for IndexedEmbeddingSourceConfig")
+        if not self.row_id_key.strip():
+            raise ValueError("row_id_key must be a non-empty string")
+
+
+class IndexedEmbeddingSource(EmbeddingArtifactSource):
+    """Eager Datarax-style source for embeddings with optional row identities."""
+
+    config: IndexedEmbeddingSourceConfig  # pyright: ignore[reportIncompatibleVariableOverride]
+
+    def __init__(
+        self,
+        config: IndexedEmbeddingSourceConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Load the embedding artifact and persist the configured row IDs."""
+        super().__init__(config, rngs=rngs, name=name)
+
+        row_id_key = config.row_id_key
+        if row_id_key is None:
+            raise ValueError("row_id_key is required for IndexedEmbeddingSource")
+
+        row_ids_array = self.artifact_metadata.get(row_id_key)
+        row_ids: tuple[str, ...] | None = None
+        if row_ids_array is not None:
+            if int(np.asarray(row_ids_array).shape[0]) != int(self.embeddings.shape[0]):
+                raise ValueError(
+                    f"{row_id_key} must contain one value per embedding row "
+                    f"({np.asarray(row_ids_array).shape[0]} vs {self.embeddings.shape[0]})."
+                )
+            row_ids = tuple(str(item) for item in np.asarray(row_ids_array))
+            cast(dict[str, Any], self.data)[row_id_key] = row_ids
+
+        object.__setattr__(self, "_row_ids", row_ids)
+
+    @property
+    def row_id_key(self) -> str:
+        """Configured metadata field name that identifies artifact rows."""
+        return cast(str, self.config.row_id_key)
+
+    @property
+    def row_ids(self) -> tuple[str, ...] | None:
+        """Tuple of persisted row identifiers, if present."""
+        return self._row_ids
+
+    def load(self) -> dict[str, Any]:
+        """Return the eager in-memory payload exposed by the source."""
+        payload = super().load()
+        if self.row_ids is not None:
+            payload[self.row_id_key] = self.row_ids
+        return payload
+
+    def align_to_reference_ids(
+        self,
+        *,
+        reference_ids: list[str] | tuple[str, ...],
+        require_row_ids: bool = True,
+        artifact_label: str,
+        id_display_name: str,
+    ) -> jnp.ndarray:
+        """Align embedding rows to a reference dataset order."""
+        canonical_reference_ids = tuple(str(item) for item in reference_ids)
+
+        if self.embeddings.ndim != 2:
+            raise ValueError(
+                f"{artifact_label} embedding artifacts must be rank-2 matrices "
+                f"(received shape {self.embeddings.shape})."
+            )
+
+        if self.row_ids is None:
+            if require_row_ids:
+                raise ValueError(
+                    f"{artifact_label} embedding artifacts must include {self.row_id_key} "
+                    "for strict alignment."
+                )
+            if self.embeddings.shape[0] != len(canonical_reference_ids):
+                raise ValueError(
+                    "Positional embedding alignment requires the same number of rows "
+                    f"as reference items ({self.embeddings.shape[0]} vs "
+                    f"{len(canonical_reference_ids)})."
+                )
+            return self.embeddings
+
+        if len(set(canonical_reference_ids)) != len(canonical_reference_ids):
+            raise ValueError(f"Reference {self.row_id_key} values must be unique.")
+        if len(set(self.row_ids)) != len(self.row_ids):
+            raise ValueError(f"Embedding artifact {self.row_id_key} values must be unique.")
+
+        reference_set = set(canonical_reference_ids)
+        artifact_set = set(self.row_ids)
+        if reference_set != artifact_set:
+            missing = sorted(reference_set - artifact_set)
+            extra = sorted(artifact_set - reference_set)
+            raise ValueError(
+                f"{id_display_name} mismatch between reference dataset and embedding artifact. "
+                f"Missing: {missing}. Extra: {extra}."
+            )
+
+        index_by_row_id = {row_id: index for index, row_id in enumerate(self.row_ids)}
+        row_indices = [index_by_row_id[row_id] for row_id in canonical_reference_ids]
+        reordered = np.asarray(self.embeddings)[row_indices]
+        return jnp.asarray(reordered, dtype=jnp.float32)

diffbio/sources/indexed_view.py

@@ -0,0 +1,191 @@
+"""IndexedViewSource - Lazy-loading view into a data source.
+
+This module provides IndexedViewSource, which wraps an existing DataSourceModule
+and provides access only to elements at specified indices. Elements are loaded
+on-demand from the underlying source, enabling lazy loading for large datasets.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Iterator
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+
+from datarax.core.config import StructuralConfig
+from datarax.core.data_source import DataSourceModule
+from datarax.typing import Element
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class IndexedViewSourceConfig(StructuralConfig):
+    """Configuration for IndexedViewSource.
+
+    Attributes:
+        shuffle: Whether to shuffle the view indices on initialization and reset
+        seed: Random seed for shuffling (optional)
+    """
+
+    shuffle: bool = False
+    seed: int | None = None
+
+
+class IndexedViewSource(DataSourceModule):
+    """Lazy-loading view into a data source using index mapping.
+
+    This source wraps an existing DataSourceModule and provides access
+    only to elements at specified indices. Elements are loaded ON-DEMAND
+    from the underlying source, enabling lazy loading for large datasets.
+
+    Key Features:
+
+    - LAZY LOADING: Elements fetched from underlying source only when accessed
+    - Memory efficient: Only stores indices, not actual data
+    - Preserves underlying source's lazy loading behavior
+    - Supports shuffling of view indices (not underlying data)
+
+    Example:
+        ```python
+        # Create view of first 1000 elements
+        indices = jnp.arange(1000)
+        config = IndexedViewSourceConfig()
+        view = IndexedViewSource(config, original_source, indices)
+        view[0]  # Fetches original_source[indices[0]] lazily
+        ```
+
+    Args:
+        config: Configuration for the view source
+        source: Underlying data source to wrap
+        indices: Array of indices into the source to expose
+        rngs: Random number generators for shuffling
+        name: Optional name for the module
+    """
+
+    def __init__(
+        self,
+        config: IndexedViewSourceConfig,
+        source: DataSourceModule,
+        indices: jnp.ndarray,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize IndexedViewSource.
+
+        Args:
+            config: Configuration for the view source
+            source: Underlying data source to wrap
+            indices: Array of indices into the source to expose
+            rngs: Random number generators for shuffling
+            name: Optional name for the module
+        """
+        super().__init__(config, rngs=rngs, name=name)
+        self._source = source
+        self._indices = indices
+        self._view_indices = jnp.arange(len(indices))  # Local view ordering
+        self._current_idx = 0
+
+        # Apply initial shuffle if configured
+        if config.shuffle:
+            self._shuffle_view()
+
+    def _shuffle_view(self) -> None:
+        """Shuffle the view indices (not the underlying data)."""
+        if self.rngs is not None and "shuffle" in self.rngs:
+            key = self.rngs.shuffle()
+        elif self.config.seed is not None:
+            key = jax.random.key(self.config.seed)
+        else:
+            key = jax.random.key(0)
+
+        self._view_indices = jax.random.permutation(key, self._view_indices)
+
+    def __len__(self) -> int:
+        """Return number of elements in the view."""
+        return len(self._indices)
+
+    def __getitem__(self, idx: int) -> Element | None:
+        """Get element at view index (LAZY - fetches from underlying source).
+
+        Args:
+            idx: Index into the VIEW (0 to len(view)-1)
+
+        Returns:
+            Element from underlying source at mapped index, or None if out of bounds
+        """
+        if idx < 0 or idx >= len(self._indices):
+            return None
+
+        # Map view index -> shuffled view index -> original source index
+        view_idx = int(self._view_indices[idx])
+        source_idx = int(self._indices[view_idx])
+
+        # LAZY: Fetch from underlying source only now
+        return self._source[source_idx]
+
+    def __iter__(self) -> Iterator[Element]:
+        """Iterate over view elements (LAZY - fetches on demand)."""
+        self._current_idx = 0
+        return self
+
+    def __next__(self) -> Element:
+        """Get next element (LAZY)."""
+        if self._current_idx >= len(self._indices):
+            raise StopIteration
+
+        element = self[self._current_idx]
+        self._current_idx += 1
+
+        if element is None:
+            raise StopIteration
+
+        return element
+
+    def reset(self, seed: int | None = None) -> None:
+        """Reset iteration and optionally reshuffle.
+
+        Args:
+            seed: Optional new seed for shuffling
+        """
+        self._current_idx = 0
+
+        if self.config.shuffle:
+            if seed is not None:
+                # Update seed and reshuffle
+                key = jax.random.key(seed)
+                self._view_indices = jax.random.permutation(key, jnp.arange(len(self._indices)))
+            else:
+                self._shuffle_view()
+
+    def get_batch(self, batch_size: int, key: jax.Array | None = None) -> list[Element]:
+        """Get next batch of elements (LAZY).
+
+        Args:
+            batch_size: Number of elements to fetch
+            key: Optional RNG key (unused, for interface compatibility)
+
+        Returns:
+            List of elements
+        """
+        batch = []
+        for _ in range(batch_size):
+            if self._current_idx >= len(self._indices):
+                break
+            element = self[self._current_idx]
+            if element is not None:
+                batch.append(element)
+            self._current_idx += 1
+        return batch
+
+    @property
+    def underlying_source(self) -> DataSourceModule:
+        """Access the underlying data source."""
+        return self._source
+
+    @property
+    def source_indices(self) -> jnp.ndarray:
+        """Get the indices into the underlying source."""
+        return self._indices