"""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']