neat\datasets.py

"""Dataset generation functions for testing BackpropNEAT."""

import numpy as np
import jax.numpy as jnp

def generate_xor_data(n_samples: int = 200, complexity: float = 1.0) -> tuple:
    """Generate complex XOR dataset with multiple clusters and rotations.
    
    Args:
        n_samples: Number of samples per quadrant
        complexity: Controls the complexity of the pattern (rotation and noise)
        
    Returns:
        Tuple of (features, labels)
    """
    points = []
    labels = []
    
    # Generate multiple clusters per quadrant
    n_clusters = 3
    samples_per_cluster = n_samples // n_clusters
    
    for cluster in range(n_clusters):
        # Add rotation to each subsequent cluster
        rotation = complexity * cluster * np.pi / 6  # 30 degree rotation per cluster
        
        # Define cluster centers with gaps
        centers = [
            # (x, y, radius, label)
            (-0.7 - 0.3*cluster, -0.7 - 0.3*cluster, 0.2, -1),  # Bottom-left
            (0.7 + 0.3*cluster, 0.7 + 0.3*cluster, 0.2, -1),   # Top-right
            (-0.7 - 0.3*cluster, 0.7 + 0.3*cluster, 0.2, 1),   # Top-left
            (0.7 + 0.3*cluster, -0.7 - 0.3*cluster, 0.2, 1),   # Bottom-right
        ]
        
        for cx, cy, radius, label in centers:
            # Generate points in a circle around center
            theta = np.random.uniform(0, 2*np.pi, samples_per_cluster)
            r = np.random.uniform(0, radius, samples_per_cluster)
            
            # Convert to cartesian coordinates
            x = r * np.cos(theta)
            y = r * np.sin(theta)
            
            # Apply rotation
            x_rot = x * np.cos(rotation) - y * np.sin(rotation)
            y_rot = x * np.sin(rotation) + y * np.cos(rotation)
            
            # Add cluster center and noise
            x = cx + x_rot + np.random.normal(0, 0.05, samples_per_cluster)
            y = cy + y_rot + np.random.normal(0, 0.05, samples_per_cluster)
            
            # Add points
            cluster_points = np.column_stack([x, y])
            points.append(cluster_points)
            labels.extend([label] * samples_per_cluster)
    
    # Convert to arrays
    X = np.vstack(points)
    y = np.array(labels, dtype=np.float32)
    
    # Add global rotation
    theta = complexity * np.pi / 4  # 45 degree global rotation
    rotation_matrix = np.array([
        [np.cos(theta), -np.sin(theta)],
        [np.sin(theta), np.cos(theta)]
    ])
    X = X @ rotation_matrix
    
    # Shuffle data
    perm = np.random.permutation(len(X))
    X = X[perm]
    y = y[perm]
    
    return jnp.array(X), jnp.array(y)

def generate_circle_data(n_samples: int = 1000, noise: float = 0.1) -> tuple:
    """Generate circle classification dataset.
    
    Args:
        n_samples: Number of samples per class
        noise: Standard deviation of Gaussian noise
        
    Returns:
        Tuple of (features, labels)
    """
    # Generate random angles
    theta = np.random.uniform(0, 2*np.pi, n_samples)
    
    # Inner circle (class -1)
    r_inner = 0.5 + np.random.normal(0, noise, n_samples)
    X_inner = np.column_stack([
        r_inner * np.cos(theta),
        r_inner * np.sin(theta)
    ])
    y_inner = np.full(n_samples, -1.0)
    
    # Outer circle (class 1)
    r_outer = 1.5 + np.random.normal(0, noise, n_samples)
    X_outer = np.column_stack([
        r_outer * np.cos(theta),
        r_outer * np.sin(theta)
    ])
    y_outer = np.full(n_samples, 1.0)
    
    # Combine and shuffle
    X = np.vstack([X_inner, X_outer])
    y = np.hstack([y_inner, y_outer])
    
    # Shuffle
    perm = np.random.permutation(len(X))
    return X[perm], y[perm]

def generate_spiral_dataset(n_points=1000, noise=0.1):
    """Generate a spiral dataset with rotation-invariant features."""
    # Generate theta values with more points near the center
    theta = np.sqrt(np.random.uniform(0, 1, n_points)) * 4 * np.pi
    
    # Generate two spirals
    data = []
    labels = []
    eps = 1e-8
    
    for i in range(n_points):
        # Base radius increases with theta
        r_base = theta[i] / (4 * np.pi)
        
        # Add noise that scales with radius
        noise_scale = noise * (1 - np.exp(-2 * r_base))
        
        for spiral_idx in range(2):
            # Rotate second spiral by pi
            angle = theta[i] + np.pi * spiral_idx
            
            # Add controlled noise to radius and angle
            r = r_base + np.random.normal(0, noise_scale)
            angle_noise = np.random.normal(0, noise_scale * 0.1)  # Less noise in angle
            angle += angle_noise
            
            # Calculate cartesian coordinates
            x = r * np.cos(angle)
            y = r * np.sin(angle)
            
            # Calculate polar coordinates
            r_point = np.sqrt(x*x + y*y)
            theta_point = np.arctan2(y, x)
            
            # Unwrap theta to handle multiple revolutions
            theta_unwrapped = theta_point + 2 * np.pi * (angle // (2 * np.pi))
            
            # Calculate spiral-specific features
            
            # 1. Local curvature (how much the spiral curves at this point)
            curvature = 1 / (r_point + eps)
            
            # 2. Spiral phase (position along spiral revolution)
            phase = theta_unwrapped % (2 * np.pi) / (2 * np.pi)
            
            # 3. Radial velocity (how fast radius changes with angle)
            dr_dtheta = 1 / (4 * np.pi)
            
            # 4. Normalized angular position (accounts for multiple revolutions)
            angular_pos = theta_unwrapped / (4 * np.pi)
            
            # 5. Spiral tightness (local measure of how tight the spiral is)
            tightness = r_point / (theta_unwrapped + eps)
            
            # 6. Relative position features (help distinguish between spirals)
            # Distance to other spiral
            other_angle = angle + np.pi
            other_x = r * np.cos(other_angle)
            other_y = r * np.sin(other_angle)
            dist_to_other = np.sqrt((x - other_x)**2 + (y - other_y)**2)
            
            # 7. Rotation-invariant features
            sin_phase = np.sin(phase * 2 * np.pi)
            cos_phase = np.cos(phase * 2 * np.pi)
            
            # Combine features with careful normalization
            features = np.array([
                x / 2.0,  # Normalize coordinates
                y / 2.0,
                r_point / 2.0,  # Normalize radius
                sin_phase,  # Already normalized
                cos_phase,  # Already normalized
                np.tanh(curvature * 2),  # Normalize curvature
                angular_pos / 2.0,  # Normalize angular position
                np.tanh(tightness),  # Normalize tightness
                np.tanh(dr_dtheta * 10),  # Normalize radial velocity
                dist_to_other / 4.0  # Normalize distance to other spiral
            ])
            
            data.append(features)
            labels.append(spiral_idx * 2 - 1)  # Convert to [-1, 1]
    
    return np.array(data), np.array(labels)

def generate_checkerboard_data(n_samples: int = 200) -> tuple:
    """Generate checkerboard dataset.
    
    Args:
        n_samples: Number of samples per class
        
    Returns:
        Tuple of (features, labels)
    """
    # Generate random points
    X = np.random.uniform(-2, 2, (n_samples * 2, 2))
    
    # Assign labels based on checkerboard pattern
    y = np.zeros(n_samples * 2)
    for i in range(len(X)):
        x1, x2 = X[i]
        y[i] = 1 if (int(np.floor(x1)) + int(np.floor(x2))) % 2 == 0 else 0
    
    return jnp.array(X), jnp.array(y)

# Export dataset functions
__all__ = ['generate_xor_data', 'generate_circle_data', 'generate_spiral_dataset', 
           'generate_checkerboard_data']