data.py

from __future__ import annotations

import numpy as np


def generate_pinwheel(samples: int, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    arms = 5
    labels = rng.integers(0, arms, size=samples)
    radius = rng.normal(1.8, 0.32, size=samples)
    tangent = rng.normal(0, 0.16, size=samples)
    base_angle = labels * 2 * np.pi / arms
    angle = base_angle + 1.25 * (radius - 1.8) + tangent
    points = np.column_stack(
        [radius * np.cos(angle), radius * np.sin(angle)]
    )
    points += rng.normal(scale=0.035, size=points.shape)
    return points.astype(np.float32)