juniper-data 0.4.2__py3-none-any.whl

juniper_data/generators/xor/generator.py

@@ -0,0 +1,162 @@
+"""Core NumPy-only XOR dataset generator.
+
+This module provides the XorGenerator class for generating XOR
+classification datasets using only NumPy operations.
+"""
+
+import numpy as np
+
+from juniper_data.core.split import shuffle_and_split
+
+from .params import XorParams
+
+VERSION = "1.0.0"
+
+
+class XorGenerator:
+    """NumPy-only generator for XOR classification datasets.
+
+    The XOR dataset consists of 4 quadrants around the origin:
+    - Quadrant 1 (++): x > 0, y > 0 -> Class 0
+    - Quadrant 2 (-+): x < 0, y > 0 -> Class 1
+    - Quadrant 3 (--): x < 0, y < 0 -> Class 0
+    - Quadrant 4 (+-): x > 0, y < 0 -> Class 1
+
+    All methods are static to ensure the generator is stateless and side-effect free.
+    """
+
+    @staticmethod
+    def generate(params: XorParams) -> dict[str, np.ndarray]:
+        """Generate a complete XOR dataset with train/test splits.
+
+        Args:
+            params: XorParams instance defining generation configuration.
+
+        Returns:
+            Dictionary containing:
+                - X_train: Training features (n_train, 2)
+                - y_train: Training labels (n_train, 2)
+                - X_test: Test features (n_test, 2)
+                - y_test: Test labels (n_test, 2)
+                - X_full: Full dataset features (total_points, 2)
+                - y_full: Full dataset labels (total_points, 2)
+        """
+        rng = np.random.default_rng(params.seed)
+
+        X, y = XorGenerator._generate_raw(params, rng)
+
+        split_result = shuffle_and_split(
+            X=X,
+            y=y,
+            train_ratio=params.train_ratio,
+            test_ratio=params.test_ratio,
+            seed=params.seed,
+            shuffle=params.shuffle,
+        )
+
+        return {
+            "X_train": split_result["X_train"],
+            "y_train": split_result["y_train"],
+            "X_test": split_result["X_test"],
+            "y_test": split_result["y_test"],
+            "X_full": X,
+            "y_full": y,
+        }
+
+    @staticmethod
+    def _generate_raw(params: XorParams, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]:
+        """Generate raw XOR coordinates and labels.
+
+        Args:
+            params: XorParams instance defining generation configuration.
+            rng: NumPy random generator for reproducibility.
+
+        Returns:
+            Tuple of (X, y) where:
+                - X: Feature array of shape (total_points, 2)
+                - y: One-hot label array of shape (total_points, 2)
+        """
+        n = params.n_points_per_quadrant
+
+        q1 = XorGenerator._generate_quadrant(
+            n_points=n,
+            x_min=params.margin,
+            x_max=params.x_range,
+            y_min=params.margin,
+            y_max=params.y_range,
+            rng=rng,
+        )
+        q2 = XorGenerator._generate_quadrant(
+            n_points=n,
+            x_min=-params.x_range,
+            x_max=-params.margin,
+            y_min=params.margin,
+            y_max=params.y_range,
+            rng=rng,
+        )
+        q3 = XorGenerator._generate_quadrant(
+            n_points=n,
+            x_min=-params.x_range,
+            x_max=-params.margin,
+            y_min=-params.y_range,
+            y_max=-params.margin,
+            rng=rng,
+        )
+        q4 = XorGenerator._generate_quadrant(
+            n_points=n,
+            x_min=params.margin,
+            x_max=params.x_range,
+            y_min=-params.y_range,
+            y_max=-params.margin,
+            rng=rng,
+        )
+
+        X = np.vstack([q1, q2, q3, q4])
+
+        if params.noise > 0:
+            X += rng.standard_normal(X.shape) * params.noise
+
+        X = X.astype(np.float32)
+
+        y = np.zeros((4 * n, 2), dtype=np.float32)
+        y[0 * n : 1 * n, 0] = 1.0
+        y[1 * n : 2 * n, 1] = 1.0
+        y[2 * n : 3 * n, 0] = 1.0
+        y[3 * n : 4 * n, 1] = 1.0
+
+        return X, y
+
+    @staticmethod
+    def _generate_quadrant(
+        n_points: int,
+        x_min: float,
+        x_max: float,
+        y_min: float,
+        y_max: float,
+        rng: np.random.Generator,
+    ) -> np.ndarray:
+        """Generate points uniformly distributed in a rectangular region.
+
+        Args:
+            n_points: Number of points to generate.
+            x_min: Minimum x value.
+            x_max: Maximum x value.
+            y_min: Minimum y value.
+            y_max: Maximum y value.
+            rng: NumPy random generator.
+
+        Returns:
+            Array of shape (n_points, 2) containing x, y coordinates.
+        """
+        x = rng.uniform(x_min, x_max, n_points)
+        y = rng.uniform(y_min, y_max, n_points)
+        return np.column_stack([x, y])
+
+
+def get_schema() -> dict:
+    """Return JSON schema describing the generator parameters.
+
+    Returns:
+        JSON schema dictionary for XorParams.
+    """
+    return XorParams.model_json_schema()

juniper_data/generators/xor/params.py

@@ -0,0 +1,30 @@
+"""Parameters for the XOR dataset generator."""
+
+from pydantic import BaseModel, Field
+
+
+class XorParams(BaseModel):
+    """Configuration parameters for XOR dataset generation.
+
+    The XOR dataset consists of 4 quadrants around the origin.
+    Points in quadrants 1 and 3 (x*y > 0) belong to class 0.
+    Points in quadrants 2 and 4 (x*y < 0) belong to class 1.
+    """
+
+    n_points_per_quadrant: int = Field(default=50, ge=1, description="Number of points per quadrant")
+    x_range: float = Field(
+        default=1.0,
+        gt=0,
+        description="Maximum absolute x value; x is sampled from the interval [-x_range, x_range]",
+    )
+    y_range: float = Field(
+        default=1.0,
+        gt=0,
+        description="Maximum absolute y value; y is sampled from the interval [-y_range, y_range]",
+    )
+    margin: float = Field(default=0.1, ge=0, description="Margin around axes (exclusion zone)")
+    noise: float = Field(default=0.0, ge=0, description="Gaussian noise level")
+    seed: int | None = Field(default=None, ge=0, description="Random seed for reproducibility")
+    train_ratio: float = Field(default=0.8, gt=0, le=1, description="Fraction of data for training")
+    test_ratio: float = Field(default=0.2, ge=0, le=1, description="Fraction of data for testing")
+    shuffle: bool = Field(default=True, description="Shuffle dataset before train/test split")

juniper_data/storage/__init__.py

@@ -0,0 +1,120 @@
+"""Storage module for dataset persistence."""
+
+from typing import TYPE_CHECKING
+
+from juniper_data.storage.base import DatasetStore
+from juniper_data.storage.cached import CachedDatasetStore
+from juniper_data.storage.local_fs import LocalFSDatasetStore
+from juniper_data.storage.memory import InMemoryDatasetStore
+
+if TYPE_CHECKING:
+    from juniper_data.storage.hf_store import HuggingFaceDatasetStore
+    from juniper_data.storage.kaggle_store import KaggleDatasetStore
+    from juniper_data.storage.postgres_store import PostgresDatasetStore
+    from juniper_data.storage.redis_store import RedisDatasetStore
+else:
+    try:
+        from juniper_data.storage.redis_store import RedisDatasetStore
+    except ImportError:
+        RedisDatasetStore = None
+
+    try:
+        from juniper_data.storage.hf_store import HuggingFaceDatasetStore
+    except ImportError:
+        HuggingFaceDatasetStore = None
+
+    try:
+        from juniper_data.storage.postgres_store import PostgresDatasetStore
+    except ImportError:
+        PostgresDatasetStore = None
+
+    try:
+        from juniper_data.storage.kaggle_store import KaggleDatasetStore
+    except ImportError:
+        KaggleDatasetStore = None
+__all__ = [
+    "DatasetStore",
+    "CachedDatasetStore",
+    "LocalFSDatasetStore",
+    "InMemoryDatasetStore",
+]
+
+if "RedisDatasetStore" in globals() and RedisDatasetStore is not None:
+    __all__.append("RedisDatasetStore")
+
+if "HuggingFaceDatasetStore" in globals() and HuggingFaceDatasetStore is not None:
+    __all__.append("HuggingFaceDatasetStore")
+
+if "PostgresDatasetStore" in globals() and PostgresDatasetStore is not None:
+    __all__.append("PostgresDatasetStore")
+
+if "KaggleDatasetStore" in globals() and KaggleDatasetStore is not None:
+    __all__.append("KaggleDatasetStore")
+
+
+def get_redis_store(**kwargs) -> "RedisDatasetStore":  # type: ignore[no-untyped-def]
+    """Get a Redis dataset store (requires redis package).
+
+    Args:
+        **kwargs: Arguments passed to RedisDatasetStore.
+
+    Returns:
+        RedisDatasetStore instance.
+
+    Raises:
+        ImportError: If redis package is not installed.
+    """
+    from juniper_data.storage.redis_store import RedisDatasetStore
+
+    return RedisDatasetStore(**kwargs)
+
+
+def get_hf_store(**kwargs) -> "HuggingFaceDatasetStore":  # type: ignore[no-untyped-def]
+    """Get a Hugging Face dataset store (requires datasets package).
+
+    Args:
+        **kwargs: Arguments passed to HuggingFaceDatasetStore.
+
+    Returns:
+        HuggingFaceDatasetStore instance.
+
+    Raises:
+        ImportError: If datasets package is not installed.
+    """
+    from juniper_data.storage.hf_store import HuggingFaceDatasetStore
+
+    return HuggingFaceDatasetStore(**kwargs)
+
+
+def get_postgres_store(**kwargs) -> "PostgresDatasetStore":  # type: ignore[no-untyped-def]
+    """Get a PostgreSQL dataset store (requires psycopg2 package).
+
+    Args:
+        **kwargs: Arguments passed to PostgresDatasetStore.
+
+    Returns:
+        PostgresDatasetStore instance.
+
+    Raises:
+        ImportError: If psycopg2 package is not installed.
+    """
+    from juniper_data.storage.postgres_store import PostgresDatasetStore
+
+    return PostgresDatasetStore(**kwargs)
+
+
+def get_kaggle_store(**kwargs) -> "KaggleDatasetStore":  # type: ignore[no-untyped-def]
+    """Get a Kaggle dataset store (requires kaggle package).
+
+    Args:
+        **kwargs: Arguments passed to KaggleDatasetStore.
+
+    Returns:
+        KaggleDatasetStore instance.
+
+    Raises:
+        ImportError: If kaggle package is not installed.
+    """
+    from juniper_data.storage.kaggle_store import KaggleDatasetStore
+
+    return KaggleDatasetStore(**kwargs)

juniper_data/storage/base.py

@@ -0,0 +1,279 @@
+"""Abstract base class for dataset storage."""
+
+from abc import ABC, abstractmethod
+
+# from collections.abc import Callable
+from datetime import UTC, datetime
+
+import numpy as np
+
+from juniper_data.core.models import DatasetMeta
+
+# from typing import Dict, List, Optional
+
+
+class DatasetStore(ABC):
+    """Abstract dataset storage interface.
+
+    Provides a common interface for storing and retrieving datasets,
+    supporting different backends (in-memory, local filesystem, cloud, etc.).
+    """
+
+    @abstractmethod
+    def save(
+        self,
+        dataset_id: str,
+        meta: DatasetMeta,
+        arrays: dict[str, np.ndarray],
+    ) -> None:
+        """Save dataset metadata and arrays.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+            meta: Dataset metadata.
+            arrays: Dictionary of numpy arrays (e.g., X_train, y_train, etc.).
+
+        Raises:
+            IOError: If the save operation fails.
+        """
+        pass
+
+    @abstractmethod
+    def get_meta(self, dataset_id: str) -> DatasetMeta | None:
+        """Get dataset metadata.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            Dataset metadata if found, None otherwise.
+        """
+        pass
+
+    @abstractmethod
+    def get_artifact_bytes(self, dataset_id: str) -> bytes | None:
+        """Get dataset artifact as bytes (NPZ format).
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            NPZ file contents as bytes if found, None otherwise.
+        """
+        pass
+
+    @abstractmethod
+    def exists(self, dataset_id: str) -> bool:
+        """Check if dataset exists.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            True if the dataset exists, False otherwise.
+        """
+        pass
+
+    @abstractmethod
+    def delete(self, dataset_id: str) -> bool:
+        """Delete dataset.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+
+        Returns:
+            True if the dataset was deleted, False if it didn't exist.
+        """
+        pass
+
+    @abstractmethod
+    def list_datasets(self, limit: int = 100, offset: int = 0) -> list[str]:
+        """List dataset IDs.
+
+        Args:
+            limit: Maximum number of dataset IDs to return.
+            offset: Number of dataset IDs to skip.
+
+        Returns:
+            List of dataset IDs.
+        """
+        pass
+
+    def update_meta(self, dataset_id: str, meta: DatasetMeta) -> bool:
+        """Update dataset metadata.
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+            meta: Updated dataset metadata.
+
+        Returns:
+            True if the dataset was updated, False if it didn't exist.
+        """
+        raise NotImplementedError("update_meta not implemented for this storage backend")
+
+    def list_all_metadata(self) -> list[DatasetMeta]:
+        """List all dataset metadata (for filtering/stats).
+
+        Returns:
+            List of all DatasetMeta objects.
+        """
+        raise NotImplementedError("list_all_metadata not implemented for this storage backend")
+
+    def record_access(self, dataset_id: str) -> None:
+        """Record an access to a dataset (updates last_accessed_at and access_count).
+
+        Args:
+            dataset_id: Unique identifier for the dataset.
+        """
+        meta = self.get_meta(dataset_id)
+        if meta is not None:
+            meta.last_accessed_at = datetime.now(UTC)
+            meta.access_count += 1
+            self.update_meta(dataset_id, meta)
+
+    def is_expired(self, meta: DatasetMeta) -> bool:
+        """Check if a dataset has expired based on its TTL.
+
+        Args:
+            meta: Dataset metadata.
+
+        Returns:
+            True if the dataset has expired, False otherwise.
+        """
+        if meta.expires_at is None:
+            return False
+        return datetime.now(UTC) > meta.expires_at
+
+    def delete_expired(self) -> list[str]:
+        """Delete all expired datasets.
+
+        Returns:
+            List of dataset IDs that were deleted.
+        """
+        deleted: list[str] = []
+        deleted.extend(
+            meta.dataset_id
+            for meta in self.list_all_metadata()
+            if self.is_expired(meta) and self.delete(meta.dataset_id)
+        )
+        return deleted
+
+    def filter_datasets(
+        self,
+        generator: str | None = None,
+        tags: list[str] | None = None,
+        tags_match: str = "any",
+        created_after: datetime | None = None,
+        created_before: datetime | None = None,
+        min_samples: int | None = None,
+        max_samples: int | None = None,
+        include_expired: bool = False,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> tuple[list[DatasetMeta], int]:
+        """Filter datasets by various criteria.
+
+        Args:
+            generator: Filter by generator name.
+            tags: Filter by tags.
+            tags_match: "any" (OR) or "all" (AND) for tag matching.
+            created_after: Filter by creation date (after).
+            created_before: Filter by creation date (before).
+            min_samples: Minimum number of samples.
+            max_samples: Maximum number of samples.
+            include_expired: Include expired datasets.
+            limit: Maximum number of results.
+            offset: Number of results to skip.
+
+        Returns:
+            Tuple of (filtered metadata list, total count before pagination).
+        """
+        all_meta = self.list_all_metadata()
+        filtered = []
+
+        for meta in all_meta:
+            if not include_expired and self.is_expired(meta):
+                continue
+            if generator is not None and meta.generator != generator:
+                continue
+            if tags is not None:
+                if tags_match == "all":
+                    if any(t not in meta.tags for t in tags):
+                        continue
+                elif all(t not in meta.tags for t in tags):
+                    continue
+            if created_after is not None and meta.created_at < created_after:
+                continue
+            if created_before is not None and meta.created_at > created_before:
+                continue
+            if min_samples is not None and meta.n_samples < min_samples:
+                continue
+            if max_samples is not None and meta.n_samples > max_samples:
+                continue
+            filtered.append(meta)
+
+        filtered.sort(key=lambda m: m.created_at, reverse=True)
+        total = len(filtered)
+        return filtered[offset : offset + limit], total
+
+    def batch_delete(self, dataset_ids: list[str]) -> tuple[list[str], list[str]]:
+        """Delete multiple datasets.
+
+        Args:
+            dataset_ids: List of dataset IDs to delete.
+
+        Returns:
+            Tuple of (deleted IDs, not found IDs).
+        """
+        deleted = []
+        not_found = []
+        for dataset_id in dataset_ids:
+            if self.delete(dataset_id):
+                deleted.append(dataset_id)
+            else:
+                not_found.append(dataset_id)
+        return deleted, not_found
+
+    def get_stats(self) -> dict[str, object]:
+        """Get aggregate statistics about stored datasets.
+
+        Returns:
+            Dictionary with statistics.
+        """
+        all_meta = self.list_all_metadata()
+
+        if not all_meta:
+            return {
+                "total_datasets": 0,
+                "total_samples": 0,
+                "by_generator": {},
+                "by_tag": {},
+                "oldest_created_at": None,
+                "newest_created_at": None,
+                "expired_count": 0,
+            }
+
+        by_generator: dict[str, int] = {}
+        by_tag: dict[str, int] = {}
+        total_samples = 0
+        expired_count = 0
+        created_times = []
+
+        for meta in all_meta:
+            by_generator[meta.generator] = by_generator.get(meta.generator, 0) + 1
+            for tag in meta.tags:
+                by_tag[tag] = by_tag.get(tag, 0) + 1
+            total_samples += meta.n_samples
+            created_times.append(meta.created_at)
+            if self.is_expired(meta):
+                expired_count += 1
+
+        return {
+            "total_datasets": len(all_meta),
+            "total_samples": total_samples,
+            "by_generator": by_generator,
+            "by_tag": by_tag,
+            "oldest_created_at": min(created_times),
+            "newest_created_at": max(created_times),
+            "expired_count": expired_count,
+        }