juniper-data 0.4.2__py3-none-any.whl

juniper_data/core/split.py

@@ -0,0 +1,120 @@
+"""Split and shuffle utilities for dataset partitioning.
+
+This module provides pure NumPy utilities for shuffling and splitting datasets
+into train/test sets with reproducible random number generation.
+"""
+
+import numpy as np
+
+
+def shuffle_data(
+    X: np.ndarray,
+    y: np.ndarray,
+    rng: np.random.Generator,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Shuffle X and y arrays together using the same permutation.
+
+    Args:
+        X: Feature array of shape (n_samples, ...).
+        y: Label array of shape (n_samples, ...).
+        rng: NumPy random generator for reproducibility.
+
+    Returns:
+        Tuple of shuffled (X, y) arrays with the same permutation applied.
+
+    Raises:
+        ValueError: If X and y have different number of samples.
+    """
+    if X.shape[0] != y.shape[0]:
+        raise ValueError(
+            f"X and y must have the same number of samples. Got X.shape[0]={X.shape[0]}, y.shape[0]={y.shape[0]}"
+        )
+
+    permutation = rng.permutation(X.shape[0])
+    return X[permutation], y[permutation]
+
+
+def split_data(
+    X: np.ndarray,
+    y: np.ndarray,
+    train_ratio: float,
+    test_ratio: float,
+) -> dict[str, np.ndarray]:
+    """Split arrays into train and test sets based on ratios.
+
+    Args:
+        X: Feature array of shape (n_samples, ...).
+        y: Label array of shape (n_samples, ...).
+        train_ratio: Fraction of data for training (0.0 to 1.0).
+        test_ratio: Fraction of data for testing (0.0 to 1.0).
+
+    Returns:
+        Dictionary with keys "X_train", "y_train", "X_test", "y_test".
+
+    Raises:
+        ValueError: If ratios are invalid or X and y have different sample counts.
+    """
+    if X.shape[0] != y.shape[0]:
+        raise ValueError(
+            f"X and y must have the same number of samples. Got X.shape[0]={X.shape[0]}, y.shape[0]={y.shape[0]}"
+        )
+
+    if not (0.0 <= train_ratio <= 1.0):
+        raise ValueError(f"train_ratio must be between 0 and 1. Got {train_ratio}")
+
+    if not (0.0 <= test_ratio <= 1.0):
+        raise ValueError(f"test_ratio must be between 0 and 1. Got {test_ratio}")
+
+    if train_ratio + test_ratio > 1.0:
+        raise ValueError(
+            f"train_ratio + test_ratio must not exceed 1.0. "
+            f"Got {train_ratio} + {test_ratio} = {train_ratio + test_ratio}"
+        )
+
+    n_samples = X.shape[0]
+    n_train = int(np.round(n_samples * train_ratio))
+    n_test = int(np.round(n_samples * test_ratio))
+
+    if n_train + n_test > n_samples:
+        n_test = n_samples - n_train
+
+    return {
+        "X_train": X[:n_train],
+        "y_train": y[:n_train],
+        "X_test": X[n_train : n_train + n_test],
+        "y_test": y[n_train : n_train + n_test],
+    }
+
+
+def shuffle_and_split(
+    X: np.ndarray,
+    y: np.ndarray,
+    train_ratio: float,
+    test_ratio: float,
+    seed: int | None = None,
+    shuffle: bool = True,
+) -> dict[str, np.ndarray]:
+    """Optionally shuffle and then split data into train/test sets.
+
+    High-level function that combines shuffling and splitting operations.
+    Uses np.random.Generator for reproducible randomness.
+
+    Args:
+        X: Feature array of shape (n_samples, ...).
+        y: Label array of shape (n_samples, ...).
+        train_ratio: Fraction of data for training (0.0 to 1.0).
+        test_ratio: Fraction of data for testing (0.0 to 1.0).
+        seed: Random seed for reproducibility. If None, uses non-deterministic seed.
+        shuffle: Whether to shuffle data before splitting. Defaults to True.
+
+    Returns:
+        Dictionary with keys "X_train", "y_train", "X_test", "y_test".
+
+    Raises:
+        ValueError: If ratios are invalid or X and y have different sample counts.
+    """
+    if shuffle:
+        rng = np.random.default_rng(seed)
+        X, y = shuffle_data(X, y, rng)
+
+    return split_data(X, y, train_ratio, test_ratio)

juniper_data/generators/__init__.py

@@ -0,0 +1,15 @@
+"""Data generators module for Juniper Data."""
+
+from .spiral import SpiralGenerator, SpiralParams
+from .spiral import get_schema as get_spiral_schema
+
+# Backwards-compatible alias: existing code may still import `get_schema`
+# from this package. Prefer `get_spiral_schema` for new code.
+get_schema = get_spiral_schema
+
+__all__ = [
+    "SpiralGenerator",
+    "SpiralParams",
+    "get_spiral_schema",
+    "get_schema",
+]

juniper_data/generators/arc_agi/__init__.py

@@ -0,0 +1,11 @@
+"""ARC-AGI (Abstraction and Reasoning Corpus) dataset generator."""
+
+from juniper_data.generators.arc_agi.generator import VERSION, ArcAgiGenerator, get_schema
+from juniper_data.generators.arc_agi.params import ArcAgiParams
+
+__all__ = [
+    "ArcAgiGenerator",
+    "ArcAgiParams",
+    "VERSION",
+    "get_schema",
+]

juniper_data/generators/arc_agi/generator.py

@@ -0,0 +1,229 @@
+"""ARC-AGI dataset generator.
+
+This module provides the ArcAgiGenerator class for loading ARC-AGI
+(Abstraction and Reasoning Corpus) tasks from Hugging Face or local files.
+"""
+
+import json
+from pathlib import Path
+
+import numpy as np
+
+from juniper_data.core.split import shuffle_and_split
+
+from .params import ArcAgiParams
+
+VERSION = "1.0.0"
+
+try:
+    from datasets import load_dataset as hf_load_dataset
+
+    HF_AVAILABLE = True
+except ImportError:
+    HF_AVAILABLE = False
+    hf_load_dataset = None  # type: ignore[assignment]
+
+
+class ArcAgiGenerator:
+    """Generator for ARC-AGI reasoning tasks.
+
+    Loads ARC tasks and converts them to padded numpy arrays suitable
+    for machine learning. Each task contains input/output grid pairs
+    demonstrating a transformation pattern.
+
+    Grid values are integers 0-9 (colors), with -1 used for padding.
+
+    Requires the `datasets` package for HuggingFace source:
+    pip install datasets
+
+    All methods are static to ensure the generator is stateless and side-effect free.
+    """
+
+    @staticmethod
+    def generate(params: ArcAgiParams) -> dict[str, np.ndarray]:
+        """Generate an ARC-AGI dataset with train/test splits.
+
+        Args:
+            params: ArcAgiParams instance defining loading configuration.
+
+        Returns:
+            Dictionary containing:
+                - X_train: Training input grids
+                - y_train: Training output grids
+                - X_test: Test input grids
+                - y_test: Test output grids
+                - X_full: All input grids
+                - y_full: All output grids
+                - task_ids: Task identifiers for each sample
+
+        Raises:
+            ImportError: If datasets package is not installed (HF source).
+            FileNotFoundError: If local path does not exist.
+        """
+        if params.source == "huggingface":
+            tasks = ArcAgiGenerator._load_from_huggingface(params)
+        else:
+            tasks = ArcAgiGenerator._load_from_local(params)
+
+        X, y, task_ids = ArcAgiGenerator._convert_tasks_to_arrays(tasks, params)
+
+        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,
+            "task_ids": task_ids,
+        }
+
+    @staticmethod
+    def _load_from_huggingface(params: ArcAgiParams) -> list[dict]:
+        """Load ARC tasks from Hugging Face Hub."""
+        if not HF_AVAILABLE:
+            raise ImportError("Hugging Face datasets package not installed. Install with: pip install datasets")
+
+        # assert hf_load_dataset is not None
+
+        try:
+            ds = hf_load_dataset("fchollet/arc-agi", split="train")  # nosec B615
+        except Exception:
+            ds = hf_load_dataset("multimodal-reasoning-lab/ARC-AGI", split="train")  # nosec B615
+
+        tasks: list[dict] = []
+        for item in ds:
+            task = {
+                "task_id": item.get("task_id", f"task_{len(tasks)}"),
+                "train": item.get("train", []),
+                "test": item.get("test", []),
+            }
+            tasks.append(task)
+
+        if params.n_tasks is not None:
+            if params.seed is None:
+                tasks = tasks[: params.n_tasks]
+
+            else:
+                rng = np.random.default_rng(params.seed)
+                indices = rng.choice(len(tasks), min(params.n_tasks, len(tasks)), replace=False)
+                tasks = [tasks[i] for i in indices]
+        return tasks
+
+    @staticmethod
+    def _load_from_local(params: ArcAgiParams) -> list[dict]:
+        """Load ARC tasks from local JSON files."""
+        if params.local_path is None:
+            raise ValueError("local_path is required when source='local'")
+
+        base_path = Path(params.local_path)
+        if not base_path.exists():
+            raise FileNotFoundError(f"Path not found: {params.local_path}")
+
+        tasks = []
+
+        if params.subset in ("training", "all"):
+            training_path = base_path / "training"
+            if training_path.exists():
+                tasks.extend(ArcAgiGenerator._load_json_dir(training_path))
+
+        if params.subset in ("evaluation", "all"):
+            eval_path = base_path / "evaluation"
+            if eval_path.exists():
+                tasks.extend(ArcAgiGenerator._load_json_dir(eval_path))
+
+        if params.n_tasks is not None:
+            if params.seed is None:
+                tasks = tasks[: params.n_tasks]
+
+            else:
+                rng = np.random.default_rng(params.seed)
+                indices = rng.choice(len(tasks), min(params.n_tasks, len(tasks)), replace=False)
+                tasks = [tasks[i] for i in indices]
+        return tasks
+
+    @staticmethod
+    def _load_json_dir(dir_path: Path) -> list[dict]:
+        """Load all JSON task files from a directory."""
+        tasks = []
+        for json_file in sorted(dir_path.glob("*.json")):
+            with open(json_file, encoding="utf-8") as f:
+                task_data = json.load(f)
+                task_data["task_id"] = json_file.stem
+                tasks.append(task_data)
+        return tasks
+
+    @staticmethod
+    def _convert_tasks_to_arrays(tasks: list[dict], params: ArcAgiParams) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Convert ARC tasks to padded numpy arrays."""
+        inputs = []
+        outputs = []
+        task_ids = []
+
+        for task in tasks:
+            task_id = task.get("task_id", "unknown")
+
+            for pair in task.get("train", []):
+                input_grid = ArcAgiGenerator._pad_grid(pair["input"], params.pad_to, params.pad_value)
+                output_grid = ArcAgiGenerator._pad_grid(pair["output"], params.pad_to, params.pad_value)
+                inputs.append(input_grid)
+                outputs.append(output_grid)
+                task_ids.append(task_id)
+
+            if params.include_test:
+                for pair in task.get("test", []):
+                    input_grid = ArcAgiGenerator._pad_grid(pair["input"], params.pad_to, params.pad_value)
+                    output_grid = ArcAgiGenerator._pad_grid(
+                        pair.get("output", [[params.pad_value]]),
+                        params.pad_to,
+                        params.pad_value,
+                    )
+                    inputs.append(input_grid)
+                    outputs.append(output_grid)
+                    task_ids.append(task_id)
+
+        if not inputs:
+            X_arr = np.zeros((0, params.pad_to * params.pad_to), dtype=np.float32)
+            y_arr = np.zeros((0, params.pad_to * params.pad_to), dtype=np.float32)
+            ids = np.array([], dtype=object)
+            return X_arr, y_arr, ids
+
+        X_stacked = np.stack(inputs)
+        y_stacked = np.stack(outputs)
+
+        if params.flatten_pairs:
+            X_arr = X_stacked.reshape(len(X_stacked), -1).astype(np.float32)
+            y_arr = y_stacked.reshape(len(y_stacked), -1).astype(np.float32)
+        else:
+            X_arr = X_stacked.astype(np.float32)
+            y_arr = y_stacked.astype(np.float32)
+
+        return X_arr, y_arr, np.array(task_ids, dtype=object)
+
+    @staticmethod
+    def _pad_grid(grid: list[list[int]], pad_to: int, pad_value: int) -> np.ndarray:
+        """Pad a grid to the specified size."""
+        arr = np.array(grid, dtype=np.int16)
+        h, w = arr.shape
+
+        padded = np.full((pad_to, pad_to), pad_value, dtype=np.int16)
+        padded[:h, :w] = arr
+
+        return padded
+
+
+def get_schema() -> dict:
+    """Return JSON schema describing the generator parameters.
+
+    Returns:
+        JSON schema dictionary for ArcAgiParams.
+    """
+    return ArcAgiParams.model_json_schema()

juniper_data/generators/arc_agi/params.py

@@ -0,0 +1,56 @@
+"""Parameters for the ARC-AGI dataset generator."""
+
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class ArcAgiParams(BaseModel):
+    """Configuration parameters for ARC-AGI dataset loading.
+
+    Loads ARC-AGI tasks from Hugging Face Hub or local JSON files.
+    The ARC (Abstraction and Reasoning Corpus) contains grid-based
+    reasoning tasks with input/output pairs.
+    """
+
+    source: Literal["huggingface", "local"] = Field(
+        default="huggingface",
+        description="Data source: 'huggingface' or 'local' JSON files",
+    )
+    local_path: str | None = Field(
+        default=None,
+        description="Path to local ARC JSON files (required if source='local')",
+    )
+    subset: Literal["training", "evaluation", "all"] = Field(
+        default="training",
+        description="Which subset to load: 'training', 'evaluation', or 'all'",
+    )
+    n_tasks: int | None = Field(
+        default=None,
+        ge=1,
+        description="Limit number of tasks to load (None for all)",
+    )
+    pad_to: int = Field(
+        default=30,
+        ge=1,
+        le=50,
+        description="Pad all grids to this size (max ARC grid is 30x30)",
+    )
+    pad_value: int = Field(
+        default=-1,
+        ge=-1,
+        le=9,
+        description="Value to use for padding (-1 recommended for masking)",
+    )
+    include_test: bool = Field(
+        default=True,
+        description="Include test input/output pairs (in addition to train pairs)",
+    )
+    flatten_pairs: bool = Field(
+        default=True,
+        description="Flatten all input/output pairs into single arrays",
+    )
+    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 before splitting")

juniper_data/generators/checkerboard/__init__.py

@@ -0,0 +1,15 @@
+"""Checkerboard classification dataset generator."""
+
+from juniper_data.generators.checkerboard.generator import (
+    VERSION,
+    CheckerboardGenerator,
+    get_schema,
+)
+from juniper_data.generators.checkerboard.params import CheckerboardParams
+
+__all__ = [
+    "CheckerboardGenerator",
+    "CheckerboardParams",
+    "VERSION",
+    "get_schema",
+]

juniper_data/generators/checkerboard/generator.py

@@ -0,0 +1,114 @@
+"""Core NumPy-only checkerboard dataset generator.
+
+This module provides the CheckerboardGenerator class for generating
+checkerboard classification datasets using only NumPy operations.
+"""
+
+import numpy as np
+
+from juniper_data.core.split import shuffle_and_split
+
+from .params import CheckerboardParams
+
+VERSION = "1.0.0"
+
+
+class CheckerboardGenerator:
+    """NumPy-only generator for checkerboard classification datasets.
+
+    Generates a 2D checkerboard pattern where alternating squares
+    belong to different classes. Points are uniformly distributed
+    across the grid.
+
+    All methods are static to ensure the generator is stateless and side-effect free.
+    """
+
+    @staticmethod
+    def generate(params: CheckerboardParams) -> dict[str, np.ndarray]:
+        """Generate a complete checkerboard dataset with train/test splits.
+
+        Args:
+            params: CheckerboardParams 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 (n_samples, 2)
+                - y_full: Full dataset labels (n_samples, 2)
+        """
+        rng = np.random.default_rng(params.seed)
+
+        X, y = CheckerboardGenerator._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: CheckerboardParams, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]:
+        """Generate raw checkerboard coordinates and labels.
+
+        Args:
+            params: CheckerboardParams instance defining generation configuration.
+            rng: NumPy random generator for reproducibility.
+
+        Returns:
+            Tuple of (X, y) where:
+                - X: Feature array of shape (n_samples, 2)
+                - y: One-hot label array of shape (n_samples, 2)
+        """
+        x_min, x_max = params.x_range
+        y_min, y_max = params.y_range
+
+        x = rng.uniform(x_min, x_max, params.n_samples)
+        y_coord = rng.uniform(y_min, y_max, params.n_samples)
+
+        X = np.column_stack([x, y_coord])
+
+        if params.noise > 0:
+            X += rng.standard_normal(X.shape) * params.noise
+
+        X = X.astype(np.float32)
+
+        x_step = (x_max - x_min) / params.n_squares
+        y_step = (y_max - y_min) / params.n_squares
+
+        x_idx = np.floor((x - x_min) / x_step).astype(int)
+        y_idx = np.floor((y_coord - y_min) / y_step).astype(int)
+
+        x_idx = np.clip(x_idx, 0, params.n_squares - 1)
+        y_idx = np.clip(y_idx, 0, params.n_squares - 1)
+
+        labels = (x_idx + y_idx) % 2
+
+        y = np.zeros((params.n_samples, 2), dtype=np.float32)
+        y[labels == 0, 0] = 1.0
+        y[labels == 1, 1] = 1.0
+
+        return X, y
+
+
+def get_schema() -> dict:
+    """Return JSON schema describing the generator parameters.
+
+    Returns:
+        JSON schema dictionary for CheckerboardParams.
+    """
+    return CheckerboardParams.model_json_schema()

juniper_data/generators/checkerboard/params.py

@@ -0,0 +1,32 @@
+"""Parameters for the checkerboard dataset generator."""
+
+from pydantic import BaseModel, Field
+
+
+class CheckerboardParams(BaseModel):
+    """Configuration parameters for checkerboard dataset generation.
+
+    Generates a checkerboard pattern classification dataset where
+    alternating squares belong to different classes.
+    """
+
+    n_samples: int = Field(default=200, ge=2, description="Total number of samples")
+    n_squares: int = Field(
+        default=4,
+        ge=2,
+        le=16,
+        description="Number of squares per side (total squares = n_squares^2)",
+    )
+    x_range: tuple[float, float] = Field(
+        default=(0.0, 1.0),
+        description="Range of x values (min, max)",
+    )
+    y_range: tuple[float, float] = Field(
+        default=(0.0, 1.0),
+        description="Range of y values (min, max)",
+    )
+    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 before splitting")

juniper_data/generators/circles/__init__.py

@@ -0,0 +1,11 @@
+"""Concentric circles classification dataset generator."""
+
+from juniper_data.generators.circles.generator import VERSION, CirclesGenerator, get_schema
+from juniper_data.generators.circles.params import CirclesParams
+
+__all__ = [
+    "CirclesGenerator",
+    "CirclesParams",
+    "VERSION",
+    "get_schema",
+]