utils.py

"""Small utility helpers shared across plot modules.

This module provides:
* ``r2_score`` — coefficient of determination for goodness-of-fit.
* ``generate_sample_linear`` — synthetic linear data for quick demos.
* ``generate_sample_nonlinear`` — synthetic nonlinear data (exp, log,
  sigmoid, power) for testing curve-fitting plots.
"""

from __future__ import annotations

import warnings

import numpy as np


def validate_arrays(
    *arrays: np.ndarray,
    names: list[str] | None = None,
) -> None:
    """Validate that arrays are 1-D, same length, non-empty, and finite.

    Parameters
    ----------
    *arrays : np.ndarray
        Arrays to validate.
    names : list[str] or None
        Descriptive names for error messages.  Defaults to ``"array_0"``, etc.

    Raises
    ------
    ValueError
        If arrays have mismatched lengths, wrong dimensionality, or are empty.
    """
    if names is None:
        names = [f"array_{i}" for i in range(len(arrays))]
    for arr, name in zip(arrays, names):
        if arr.ndim != 1:
            raise ValueError(f"{name} must be 1-D, got {arr.ndim}-D")
        if len(arr) == 0:
            raise ValueError(f"{name} is empty")
        n_bad = int(np.sum(~np.isfinite(arr)))
        if n_bad:
            warnings.warn(f"{name} contains {n_bad} non-finite value(s)", stacklevel=3)
    if len(arrays) > 1:
        ref = len(arrays[0])
        for arr, name in zip(arrays[1:], names[1:]):
            if len(arr) != ref:
                raise ValueError(
                    f"Length mismatch: {names[0]} has {ref} elements, "
                    f"{name} has {len(arr)}"
                )


def validate_matrix(
    data: np.ndarray,
    name: str = "data",
) -> None:
    """Validate that *data* is a non-empty 2-D array.

    Parameters
    ----------
    data : np.ndarray
        Matrix to validate.
    name : str
        Descriptive name for error messages.

    Raises
    ------
    ValueError
        If *data* is not 2-D or is empty.
    """
    if data.ndim != 2:
        raise ValueError(f"{name} must be 2-D, got {data.ndim}-D")
    if data.size == 0:
        raise ValueError(f"{name} is empty")


def r2_score(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """Coefficient of determination (R²).

    Parameters
    ----------
    y_true : array
        Ground-truth (observed) values.
    y_pred : array
        Predicted values from a model.

    Returns
    -------
    float
        R² score.  1.0 = perfect fit; 0.0 = no better than mean;
        negative = worse than predicting the mean.

    Example
    -------
    >>> y_true = np.array([1, 2, 3])
    >>> y_pred = np.array([1.1, 2.0, 2.9])
    >>> r2_score(y_true, y_pred)
    0.995
    """
    ss_res = np.sum((y_true - y_pred) ** 2)
    ss_tot = np.sum((y_true - np.mean(y_true)) ** 2)
    return float(1.0 - ss_res / ss_tot)


def generate_sample_linear(
    n: int = 30,
    slope: float = 2.0,
    intercept: float = 1.0,
    noise: float = 0.8,
    seed: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """Generate synthetic linear data ``y = slope * x + intercept + noise``.

    Parameters
    ----------
    n : int
        Number of data points.
    slope : float
        True slope of the linear relationship.
    intercept : float
        True y-intercept.
    noise : float
        Standard deviation of Gaussian noise added to y.
    seed : int
        Random seed for reproducibility.

    Returns
    -------
    tuple[np.ndarray, np.ndarray]
        ``(x, y)`` arrays sorted by x.
    """
    rng = np.random.default_rng(seed)
    x = np.sort(rng.uniform(0, 10, n))
    y = slope * x + intercept + rng.normal(0, noise, n)
    return x, y


def generate_sample_nonlinear(
    kind: str = "exp",
    n: int = 40,
    noise: float = 0.3,
    seed: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """Generate synthetic nonlinear data for testing curve-fitting plots.

    Parameters
    ----------
    kind : str
        Type of relationship.  One of:
        ``"exp"``     — exponential:  ``y = 2·exp(0.5·x) + noise``
        ``"log"``     — logarithmic:  ``y = 3·ln(x) + 1 + noise``
        ``"sigmoid"`` — logistic:     ``y = 6 / (1+exp(-2(x-2.5))) + noise``
        ``"power"``   — power law:    ``y = 1.5·x^1.8 + noise``
    n : int
        Number of data points.
    noise : float
        Standard deviation of Gaussian noise.
    seed : int
        Random seed for reproducibility.

    Returns
    -------
    tuple[np.ndarray, np.ndarray]
        ``(x, y)`` arrays sorted by x.

    Raises
    ------
    ValueError
        If *kind* is not one of the supported strings.
    """
    rng = np.random.default_rng(seed)
    x = np.sort(rng.uniform(0.5, 5, n))
    match kind:
        case "exp":
            y = 2.0 * np.exp(0.5 * x) + rng.normal(0, noise, n)
        case "log":
            y = 3.0 * np.log(x) + 1.0 + rng.normal(0, noise, n)
        case "sigmoid":
            y = 6.0 / (1.0 + np.exp(-2.0 * (x - 2.5))) + rng.normal(0, noise, n)
        case "power":
            y = 1.5 * x**1.8 + rng.normal(0, noise, n)
        case _:
            raise ValueError(f"Unknown kind: {kind!r}")
    return x, y