derivkit 1.0.0__py3-none-any.whl

derivkit/likelihoods/poisson.py

@@ -0,0 +1,176 @@
+"""Poissonian likelihoods function module."""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.stats import poisson
+
+__all__ = [
+    "build_poissonian_likelihood",
+]
+
+
+def build_poissonian_likelihood(
+    data: float | np.ndarray[float],
+    model_parameters: float | np.ndarray[float],
+    return_log: bool = True,
+    ) -> tuple[np.ndarray[float], np.ndarray[float]]:
+    """Constructs the Poissonian likelihoods function.
+
+    The shape of the data products depend on the shape of ``model_parameters``.
+    The assumption is that ``model_parameters`` contains the expectation value
+    of some quantity which is either uniform for the entire distribution or is
+    distributed across a grid of bins. It is uniform for the entire distribution
+    if it is a scalar.
+
+    The function will try to reshape ``data`` to align with ``model_parameters``.
+    If ``model_parameters`` is a scalar, then ``data`` will be flattened. Otherwise,
+    the grid can contain any number of axes, but currently the number of axes
+    is hardcoded to 2. Supplying a higher-dimensional array to
+    ``model_parameters`` may produce unexpected results.
+
+    This hardcoded limit means that, while it is possible to supply
+    ``model_parameters`` along a 1D grid, the output shape will always be a
+    2D row-major array. See Examples for more details.
+
+    Args:
+        data: an array representing the given data values.
+        model_parameters: an array representing the means of the data samples.
+        return_log: when set to ``True``, returns the log-likelihoods instead of
+                the probability mass function.
+
+    Returns:
+        A tuple of arrays containing (in order):
+
+            - the data, reshaped to align with the model parameters.
+            - the values of the Poissonian probability mass function computed
+              from the data and model parameters.
+
+    Raises:
+        ValueError: If any of the model_parameters are negative or non-finite,
+            or the data points cannot be reshaped to align with
+            model_parameters.
+
+    Examples:
+        Scalar mean + scalar data:
+
+        >>> import numpy as np
+        >>> from scipy.stats import poisson
+        >>> from derivkit.likelihoods.poisson import build_poissonian_likelihood
+        >>> x, y = build_poissonian_likelihood(2, 1.4, return_log=False)
+        >>> x.shape, y.shape
+        ((1,), (1,))
+        >>> x[0].item()
+        2
+        >>> np.allclose(y[0], poisson.pmf(2, 1.4))
+        True
+
+        Vector data + scalar mean (data are flattened):
+
+        >>> data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
+        >>> model_parameters = 2.4
+        >>> x, y = build_poissonian_likelihood(
+        ...             data, model_parameters, return_log=False
+        ... )
+        >>> x.shape, y.shape
+        ((10,), (10,))
+        >>> np.array_equal(x, data)
+        True
+        >>> np.allclose(y, poisson.pmf(data, 2.4))
+        True
+
+        Shape follows ``model_parameters``:
+
+        >>> data = np.array([1, 2])
+        >>> model_parameters = np.array([3])
+        >>> x, y = build_poissonian_likelihood(
+        ...             data, model_parameters, return_log=False
+        ... )
+        >>> x.shape, y.shape
+        ((2, 1), (2, 1))
+        >>> np.array_equal(x[:, 0], data)
+        True
+        >>> np.allclose(y[:, 0], poisson.pmf(data, 3))
+        True
+
+        1D grid of bins produces a row-major 2D output:
+
+        >>> model_parameters = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
+        >>> data = np.array([1, 2, 3, 4, 5, 6])
+        >>> x, y = build_poissonian_likelihood(
+        ...             data, model_parameters, return_log=False
+        ... )
+        >>> x.shape, y.shape
+        ((1, 6), (1, 6))
+        >>> np.array_equal(x[0], data)
+        True
+        >>> np.allclose(y[0], poisson.pmf(data, model_parameters))
+        True
+
+        2D grid:
+
+        >>> data = np.array([[1, 2, 3], [4, 5, 6]])
+        >>> model_parameters = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
+        >>> x, y = build_poissonian_likelihood(
+        ...             data, model_parameters, return_log=False
+        ... )
+        >>> x.shape, y.shape
+        ((1, 2, 3), (1, 2, 3))
+        >>> np.array_equal(x[0], data)
+        True
+        >>> np.allclose(y[0], poisson.pmf(data, model_parameters))
+        True
+
+        Stacked data on the same grid:
+
+        >>> val1 = np.array([[1, 2, 3], [4, 5, 6]])
+        >>> val2 = np.array([[7, 8, 9], [10, 11, 12]])
+        >>> data = np.array([val1, val2])
+        >>> model_parameters = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
+        >>> x, y = build_poissonian_likelihood(
+        ...             data, model_parameters, return_log=False
+        ... )
+        >>> x.shape, y.shape
+        ((2, 2, 3), (2, 2, 3))
+        >>> np.array_equal(x[0], val1) and np.array_equal(x[1], val2)
+        True
+        >>> np.allclose(y[0], poisson.pmf(val1, model_parameters))
+        True
+        >>> np.allclose(y[1], poisson.pmf(val2, model_parameters))
+        True
+
+        Same result when supplying flattened data:
+
+        >>> data_flat = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
+        >>> x2, y2 = build_poissonian_likelihood(
+        ...             data_flat, model_parameters, return_log=False
+        ... )
+        >>> np.array_equal(x2, x) and np.allclose(y2, y)
+        True
+    """
+    values_to_reshape = np.asarray(data)
+    parameters = np.asarray(model_parameters)
+
+    if np.any(values_to_reshape < 0):
+        raise ValueError("values of data must be non-negative.")
+    if np.any(~np.isfinite(values_to_reshape)):
+        raise ValueError("values of data must be finite.")
+    if np.any(parameters < 0):
+        raise ValueError("values of model_parameters must be non-negative.")
+    if np.any(~np.isfinite(parameters)):
+        raise ValueError("values of model_parameters must be finite.")
+
+    try:
+        counts = values_to_reshape.reshape(-1, *parameters.shape[-2:])
+    except ValueError:
+        raise ValueError(
+            "data cannot be reshaped to align with model_parameters: "
+            f"data.shape={values_to_reshape.shape} is incompatible with "
+            f"model_parameters.shape={parameters.shape}."
+        )
+
+    probabilities = poisson.logpmf(counts, parameters) \
+        if return_log \
+        else poisson.pmf(counts, parameters)
+
+    return counts, probabilities

derivkit/utils/__init__.py

@@ -0,0 +1,13 @@
+"""Utility functions for DerivKit package."""
+
+from derivkit.utils.linalg import (
+    invert_covariance,
+    normalize_covariance,
+    solve_or_pinv,
+)
+
+__all__ = [
+    "solve_or_pinv",
+    "invert_covariance",
+    "normalize_covariance",
+]

derivkit/utils/concurrency.py

@@ -0,0 +1,213 @@
+"""Concurrency management for derivative computations."""
+
+from __future__ import annotations
+
+import contextvars
+import os
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from typing import Any, Callable, Iterator, Sequence, Tuple
+
+__all__ = [
+    "set_default_inner_derivative_workers",
+    "set_inner_derivative_workers",
+    "resolve_inner_from_outer",
+    "parallel_execute",
+    "_inner_workers_var",
+    "normalize_workers",
+    "resolve_workers",
+]
+
+
+# Context-var and default
+_inner_workers_var: contextvars.ContextVar[int | None] = contextvars.ContextVar(
+    "derivkit_inner_workers", default=None
+)
+_DEFAULT_INNER_WORKERS: int | None = None
+
+
+def set_default_inner_derivative_workers(n: int | None) -> None:
+    """Sets the module-wide default for inner derivative workers.
+
+    Args:
+        n: Number of inner derivative workers, or None for automatic policy.
+
+    Returns:
+        None
+    """
+    global _DEFAULT_INNER_WORKERS
+    _DEFAULT_INNER_WORKERS = None if n is None else int(n)
+
+
+@contextmanager
+def set_inner_derivative_workers(n: int | None) -> Iterator[int | None]:
+    """Temporarily sets the number of inner derivative workers.
+
+    Args:
+        n: Number of inner derivative workers, or ``None`` for automatic policy.
+
+    Yields:
+        int | None: The previous worker setting (restored on exit).
+    """
+    prev = _inner_workers_var.get()
+    token = _inner_workers_var.set(None if n is None else int(n))
+    try:
+        yield prev
+    finally:
+        _inner_workers_var.reset(token)
+
+
+def _int_env(name: str) -> int | None:
+    """Reads a positive integer from an environment variable, or None if unset/invalid.
+
+    Args:
+        name: Environment variable name.
+
+    Returns:
+        Positive integer value, or None.
+    """
+    v = os.getenv(name)
+    if not v:
+        return None
+    try:
+        i = int(v)
+        return i if i > 0 else None
+    except ValueError:
+        return None
+
+def _detect_hw_threads() -> int:
+    """Detects the number of hardware threads, capped by relevant environment variables.
+
+    Returns:
+        Number of hardware threads (at least 1).
+    """
+    hints = [
+        _int_env("OMP_NUM_THREADS"),
+        _int_env("MKL_NUM_THREADS"),
+        _int_env("OPENBLAS_NUM_THREADS"),
+        _int_env("VECLIB_MAXIMUM_THREADS"),
+        _int_env("NUMEXPR_NUM_THREADS"),
+    ]
+    env_cap = min([h for h in hints if h is not None], default=None)
+    hw = os.cpu_count() or 1
+    return max(1, min(hw, env_cap) if env_cap else hw)
+
+
+def resolve_inner_from_outer(w_params: int) -> int | None:
+    """Resolves the number of inner derivative workers based on outer workers and defaults.
+
+    Args:
+        w_params: Number of outer derivative workers.
+
+    Returns:
+        Number of inner derivative workers, or None for automatic policy.
+    """
+    w = _inner_workers_var.get()
+    if w is not None:
+        return w
+    if _DEFAULT_INNER_WORKERS is not None:
+        return _DEFAULT_INNER_WORKERS
+    cores = _detect_hw_threads()
+    if w_params > 1:
+        return min(4, max(1, cores // w_params))
+    return min(4, cores)
+
+
+def parallel_execute(
+    worker: Callable[..., Any],
+    arg_tuples: Sequence[Tuple[Any, ...]],
+    *,
+    outer_workers: int = 1,
+    inner_workers: int | None = None,
+    backend: str = "threads",
+) -> list[Any]:
+    """Applies a function to groups of arguments in parallel.
+
+    Inner worker setting is applied to the context, so calls inside worker
+    will see the resolved inner worker count.
+
+    Args:
+        worker: Function applied to each entry in ``arg_tuples`` (called as ``worker(*args)``).
+        arg_tuples: Argument tuples; each tuple is expanded into one ``worker(*args)`` call.
+        outer_workers: Parallelism level for outer execution.
+        inner_workers: Inner derivative worker setting to propagate via contextvar.
+        backend: Parallel backend. Currently supported: "threads".
+
+    Returns:
+        List of worker return values.
+    """
+    backend_l = str(backend).lower()
+    if backend_l not in {"threads"}:
+        raise NotImplementedError(
+            f"parallel_execute backend={backend!r} not supported yet."
+            f" Use backend='threads'."
+        )
+
+    with set_inner_derivative_workers(inner_workers):
+        if outer_workers > 1:
+            with ThreadPoolExecutor(max_workers=outer_workers) as ex:
+                futures = []
+                for args in arg_tuples:
+                    # Each task gets its own copy of the current context
+                    ctx = contextvars.copy_context()
+                    futures.append(ex.submit(ctx.run, worker, *args))
+                return [f.result() for f in futures]
+        else:
+            return [worker(*args) for args in arg_tuples]
+
+
+def normalize_workers(
+    n_workers: Any
+) -> int:
+    """Ensures n_workers is a positive integer, defaulting to 1.
+
+    Args:
+        n_workers: Input number of workers (can be None, float, negative, etc.)
+
+    Returns:
+        int: A positive integer number of workers (at least 1).
+
+    Raises:
+        None: Invalid inputs are coerced to 1.
+    """
+    try:
+        n = int(n_workers)
+    except (TypeError, ValueError):
+        n = 1
+    return 1 if n < 1 else n
+
+
+def resolve_workers(
+    n_workers: Any,
+    dk_kwargs: dict[str, Any],
+) -> tuple[int, int | None, dict[str, Any]]:
+    """Decides how parallel work is split between outer calculus routines and the inner derivative engine.
+
+    Outer workers parallelize across independent derivative tasks (e.g. parameters,
+    output components, Hessian entries). Inner workers control parallelism inside
+    each derivative evaluation (within DerivativeKit).
+
+    If both levels spawn workers simultaneously, nested parallelism can cause
+    oversubscription. By default, the inner worker count is derived from the
+    outer worker count to avoid that. You can override this by passing
+    ``inner_workers=<int>`` via ``dk_kwargs``.
+
+    Args:
+        n_workers: Number of outer workers. If ``None``, defaults to 1.
+        dk_kwargs: Keyword arguments forwarded to DerivativeKit.differentiate.
+            May include ``inner_workers`` to override the default policy.
+
+    Returns:
+        (outer_workers, inner_workers, dk_kwargs_cleaned), where ``dk_kwargs_cleaned``
+        has any ``inner_workers`` entry removed.
+    """
+    dk_kwargs_cleaned = dict(dk_kwargs)
+    inner_override = dk_kwargs_cleaned.pop("inner_workers", None)
+
+    outer = normalize_workers(n_workers)
+    if inner_override is None:
+        inner = resolve_inner_from_outer(outer)
+    else:
+        inner = normalize_workers(inner_override)
+
+    return outer, inner, dk_kwargs_cleaned

derivkit/utils/extrapolation.py

@@ -0,0 +1,254 @@
+"""Extrapolation methods for numerical approximations."""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+import numpy as np
+from numpy.typing import NDArray
+
+__all__ = [
+    "richardson_extrapolate",
+    "ridders_extrapolate",
+    "gauss_richardson_extrapolate",
+]
+
+
+def richardson_extrapolate(
+        base_values: Sequence[NDArray[np.float64] | float],
+        p: int,
+        r: float = 2.0,
+) -> NDArray[np.float64] | float:
+    """Computes Richardson extrapolation on a sequence of approximations.
+
+    Richardson extrapolation improves the accuracy of a sequence of
+    numerical approximations that converge with a known leading-order error
+    term. Given a sequence of approximations computed with decreasing step sizes,
+    this method combines them to eliminate the leading error term, yielding
+    a more accurate estimate of the true value.
+
+    Args:
+        base_values:
+            Sequence of approximations at different step sizes.
+            The step sizes are assumed to decrease by a factor of ``r``
+            between successive entries.
+        p:
+            The order of the leading error term in the approximations.
+        r:
+            The step-size reduction factor between successive entries
+            (default is ``2.0``).
+
+    Returns:
+        The extrapolated value with improved accuracy.
+
+    Raises:
+        ValueError: If ``base_values`` has fewer than two entries.
+    """
+    # Work on float arrays for both scalar and vector cases
+    n = len(base_values)
+    if n < 2:
+        raise ValueError("richardson_extrapolate requires at least two base values.")
+
+    vals = [np.asarray(v, dtype=float) for v in base_values]
+
+    for j in range(1, n):
+        factor = r ** (p * j)
+        for k in range(n - 1, j - 1, -1):
+            vals[k] = (factor * vals[k] - vals[k - 1]) / (factor - 1.0)
+
+    result = vals[-1]
+    return float(result) if result.ndim == 0 else result
+
+
+def ridders_extrapolate(
+    base_values: Sequence[NDArray[np.float64] | float],
+    r: float = 2.0,
+    *,
+    extrapolator = richardson_extrapolate,
+    p: int = 2,
+) -> tuple[NDArray[np.float64] | float, float]:
+    """Computes a Ridders-style extrapolation on a sequence of approximations.
+
+    This builds the usual Ridders diagonal assuming a central finite-difference
+    scheme (leading error is approximately O(h^2)) by repeatedly extrapolating
+    prefixes of ``base_values``. By default it uses
+    :func:`derivkit.utils.extrapolation.richardson_extrapolate`
+    with ``p=2``, but a different extrapolator can be passed if needed.
+
+    Args:
+        base_values:
+            Sequence of derivative approximations at step sizes
+            h, h/r, h/r^2, ... (all same shape: scalar, vector, or tensor).
+        r:
+            Step-size reduction factor (default ``2.0``).
+        extrapolator:
+            Function implementing the extrapolation step. Must have the
+            signature ``extrapolator(base_values, p, r) -> array_like``.
+            Defaults to
+            :func:`derivkit.utils.extrapolation.richardson_extrapolate`.
+        p:
+            Leading error order passed to ``extrapolator`` (default ``2``).
+
+    Returns:
+        A tuple ``(best_value, error_estimate)`` where:
+
+        * ``best_value`` is the extrapolated estimate chosen from the
+          diagonal entries.
+        * ``error_estimate`` is a heuristic scalar error scale given by the
+          minimum difference between consecutive diagonal elements.
+
+    Raises:
+        ValueError:
+            If fewer than two base values are provided.
+    """
+    n = len(base_values)
+    if n < 2:
+        raise ValueError("ridders_extrapolate requires at least two base values.")
+
+    diag: list[NDArray[np.float64]] = []
+    err_estimates: list[float] = []
+
+    for j in range(n):
+        if j == 0:
+            d_j = np.asarray(base_values[0], dtype=float)
+        else:
+            # Use the chosen extrapolator on the first (j+1) base values
+            d_j = np.asarray(
+                extrapolator(base_values[: j + 1], p=p, r=r),
+                dtype=float,
+            )
+
+        diag.append(d_j)
+
+        if j == 0:
+            err_estimates.append(np.inf)
+        else:
+            diff = np.asarray(diag[j] - diag[j - 1], dtype=float)
+            err_estimates.append(float(np.max(np.abs(diff))))
+
+    # Pick the diagonal element with the smallest estimated error
+    best_idx = int(np.argmin(err_estimates))
+    best_val = diag[best_idx]
+    best_err = err_estimates[best_idx]
+
+    if best_val.ndim == 0:
+        return float(best_val), float(best_err)
+    return best_val, float(best_err)
+
+
+def _rbf_kernel_1d(x: NDArray[np.float64],
+                   y: NDArray[np.float64],
+                   length_scale: float) -> NDArray[np.float64]:
+    """Compute the RBF kernel matrix between 1D inputs x and y.
+
+    Args:
+        x: 1D array of shape (n,).
+        y: 1D array of shape (m,).
+        length_scale: Length scale parameter for the RBF kernel.
+
+    Returns:
+        Kernel matrix of shape (n, m).
+    """
+    x = np.atleast_1d(x).astype(float)
+    y = np.atleast_1d(y).astype(float)
+    diff2 = (x[:, None] - y[None, :]) ** 2
+    return np.exp(-0.5 * diff2 / (length_scale**2))
+
+
+def gauss_richardson_extrapolate(
+    base_values: Sequence[NDArray[np.float64] | float],
+    h_values: Sequence[float],
+    p: int,
+    jitter: float = 1e-10,
+) -> tuple[NDArray[np.float64] | float, NDArray[np.float64] | float]:
+    """Gauss–Richardson extrapolation for a sequence of approximations f(h_i).
+
+    This method uses a Gaussian-process model with a radial-basis-function (RBF)
+    kernel to perform Richardson extrapolation, providing both an improved estimate
+    of the true value at h=0 and an uncertainty estimate. For more details, see arXiv:2401.07562.
+
+    Args:
+        base_values: Sequence of approximations at different step sizes h_i.
+        h_values: Corresponding step sizes (must be positive and same length as base_values).
+        p: The order of the leading error term in the approximations.
+        jitter: Small positive value added to the diagonal of the kernel matrix for numerical stability.
+            Defaults to ``1e-10``.
+
+    Returns:
+        A tuple (extrapolated_value, error_estimate) where:
+
+          - extrapolated_value is the Gauss–Richardson extrapolated estimate at h=0.
+          - error_estimate is a heuristic uncertainty estimate for the extrapolated value.
+
+    Raises:
+        ValueError:
+            If h_values and base_values have different lengths or if any h_value is non-positive.
+    """
+    h = np.asarray(h_values, dtype=float).ravel()
+    if len(base_values) != h.size:
+        raise ValueError("base_values and h_values must have the same length.")
+    if np.any(h <= 0):
+        raise ValueError("All h_values must be > 0.")
+
+    y = np.stack([np.asarray(v, dtype=float) for v in base_values], axis=0)
+    n = h.size
+
+    # Error bound b(h) = h^p
+    b = h**p
+
+    # crude length scale from spacing
+    h_sorted = np.sort(h)
+    # then we compute the differences between consecutive sorted h values
+    diffs = np.diff(h_sorted)
+    # if there are any positive differences, take the median of those
+    if np.any(diffs > 0):
+        char = np.median(diffs[diffs > 0])
+    else:
+        char = max(h.max() - h.min(), 1e-12)
+
+    ell = char
+
+    # we then build the kernel matrix kb
+    ke = _rbf_kernel_1d(h, h, ell)
+    kb = (b[:, None] * b[None, :]) * ke
+    kb += jitter * np.eye(n)
+
+    # we then precompute the matrix-vector product kb^{-1} 1
+    one = np.ones(n)
+    kb_inv_1 = np.linalg.solve(kb, one)
+
+    flat = y.reshape(n, -1)
+    means = []
+    errs = []
+
+    denom = float(one @ kb_inv_1)
+    for j in range(flat.shape[1]):
+        col = flat[:, j]
+
+        # reuse kb_inv_1 or recompute:
+        kb_inv_y = np.linalg.solve(kb, col)
+
+        num = float(one @ kb_inv_y)
+        mean0 = num / denom  # μ̂
+
+        # Residuals
+        resid = col - mean0 * one
+        kb_inv_resid = np.linalg.solve(kb, resid)
+
+        # Noise variance estimate
+        sigma2 = float(resid @ kb_inv_resid) / max(n - 1, 1)
+
+        # Variance at h=0
+        var0 = sigma2 / denom if denom > 0 else 0.0
+        var0 = max(var0, 0.0)
+        std0 = float(np.sqrt(var0))
+
+        means.append(mean0)
+        errs.append(std0)
+
+    means_arr = np.array(means).reshape(y.shape[1:])
+    errs_arr = np.array(errs).reshape(y.shape[1:])
+
+    if means_arr.ndim == 0:
+        return float(means_arr), float(errs_arr)
+    return means_arr, errs_arr