derivkit 1.0.0__py3-none-any.whl

derivkit/derivatives/adaptive/transforms.py

@@ -0,0 +1,245 @@
+"""Helpers for parameter transformations and converting derivatives between coordinate systems.
+
+This module provides small, self-contained transforms that make adaptive
+polynomial fitting robust near parameter boundaries.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import numpy as np
+
+__all__ = [
+    "signed_log_forward",
+    "signed_log_to_physical",
+    "signed_log_derivatives_to_x",
+    "sqrt_domain_forward",
+    "sqrt_to_physical",
+    "sqrt_derivatives_to_x_at_zero",
+]
+
+
+def signed_log_forward(x0: float) -> Tuple[float, float]:
+    """Computes the signed-log coordinates for an expansion point.
+
+    The *signed-log* map represents a physical coordinate ``x`` as
+    ``x = sgn * exp(q)``, where ``q = log(|x|)`` and ``sgn = sign(x)``.
+    Here, **physical** means the model’s native parameter (``x``), while
+    **internal** means the reparameterized coordinate used for numerics (``q``).
+    This reparameterization keeps multiplicative variation (orders of magnitude)
+    well-behaved and avoids crossing through zero during local polynomial fits.
+
+    Args:
+        x0: Expansion point in physical coordinates. Must be finite and non-zero.
+
+    Returns:
+        Tuple[float, float]: ``(q0, sgn)``, where ``q0 = log(|x0|)`` and
+        ``sgn = +1.0`` if ``x0 > 0`` else ``-1.0``.
+
+    Raises:
+        ValueError: If ``x0`` is not finite or equals zero.
+    """
+    if not np.isfinite(x0):
+        raise ValueError("signed_log_forward requires a finite value of x0.")
+    if x0 == 0.0:
+        raise ValueError("signed_log_forward requires that x0 is non-zero.")
+    sgn = 1.0 if x0 > 0.0 else -1.0
+    q0 = np.log(abs(x0))
+    return q0, sgn
+
+
+def signed_log_to_physical(q: np.ndarray, sgn: float) -> np.ndarray:
+    """Maps internal signed-log coordinate(s) to physical coordinate(s).
+
+    Args:
+        q: Internal coordinate(s) q = log(abs(x)).
+        sgn: Fixed sign (+1 or -1) taken from sign(x0).
+
+    Returns:
+        Physical coordinate(s) x = sgn * exp(q).
+
+    Raises:
+        ValueError: If `sgn` is not +1 or -1, or if `q` contains non-finite values.
+    """
+    try:
+        sgn = _normalize_sign(sgn)
+    except ValueError as e:
+        raise ValueError(f"signed_log_to_physical: invalid `sgn`: {e}") from None
+
+    q = np.asarray(q, dtype=float)
+    try:
+        _require_finite("q", q)
+    except ValueError as e:
+        raise ValueError(f"signed_log_to_physical: {e}") from None
+
+    return sgn * np.exp(q)
+
+
+def signed_log_derivatives_to_x(
+    order: int,
+    x0: float,
+    dfdq: np.ndarray,
+    d2fdq2: Optional[np.ndarray] = None,
+) -> np.ndarray:
+    """Converts derivatives from the signed-log coordinate ``q`` to the original parameter ``x`` at ``x0 ≠ 0``.
+
+    This method uses the chain rule to convert derivatives computed in the
+    internal signed-log coordinate q back to physical coordinates x at a
+    non-zero expansion point x0.
+
+    Args:
+        order: Derivative order to return (1 or 2).
+        x0: Expansion point in the original parameter ``x`` (the model’s native coordinate);
+            must be finite and non-zero.
+        dfdq: First derivative in q (shape: (n_components,) or broadcastable).
+        d2fdq2: Second derivative with respect to ``q``; required when ``order == 2``.
+            A 1-D array with one value per component (shape ``(n_components,)``) or
+            broadcastable to that.
+
+    Returns:
+        The derivative(s) in physical coordinates at x0.
+
+    Raises:
+        ValueError: If `x0 == 0`, if required inputs (d2fdq2) are missing for order=2,
+            or if `x0` is not finite.
+        NotImplementedError: If `order` not in {1, 2}.
+    """
+    if not np.isfinite(x0) or x0 == 0.0:
+        raise ValueError("signed_log_derivatives_to_x requires finite x0 != 0.")
+    dfdq = np.asarray(dfdq, dtype=float)
+    if order == 1:
+        return dfdq / x0
+    elif order == 2:
+        if d2fdq2 is None:
+            raise ValueError("order=2 conversion requires d2fdq2.")
+        d2fdq2 = np.asarray(d2fdq2, dtype=float)
+        return (d2fdq2 - dfdq) / (x0 ** 2)
+    raise NotImplementedError("signed_log_derivatives_to_x supports orders 1 and 2.")
+
+
+def sqrt_domain_forward(x0: float) -> tuple[float, float]:
+    """Computes the internal domain coordinate u0 for the square-root domain transformation.
+
+    The *square-root domain* transform re-expresses a parameter ``x`` as
+    ``x = s * u**2``, where ``s`` is the domain sign (+1 or –1).  This mapping
+    flattens steep behavior near a boundary such as ``x = 0`` and allows
+    smooth polynomial fitting on either the positive or negative side.
+
+    Args:
+        x0: Expansion point in physical coordinates (finite). May be ±0.0
+
+    Returns:
+        Tuple[float, float]: ``(u0, s)``, with u0 >= 0 and sgn in {+1.0, -1.0}.
+    """
+    if not np.isfinite(x0):
+        raise ValueError("sqrt_domain_forward requires finite x0.")
+    sgn = _sgn_from_x0(x0)
+    u0 = 0.0 if x0 == 0.0 else float(np.sqrt(abs(x0)))
+    return u0, sgn
+
+
+def sqrt_to_physical(u: np.ndarray, sign: float) -> np.ndarray:
+    """Maps internal domain coordinate(s) to physical coordinate(s).
+
+    This method maps internal coordinate(s) u to physical coordinate(s) x
+    using the relation x = sign * u^2.
+
+    Args:
+        u: Internal coordinate(s).
+        sign: Domain sign (+1 for x ≥ 0, -1 for x ≤ 0).
+
+    Returns:
+        Physical coordinate(s) x = sign * u^2.
+
+    Raises:
+        ValueError: If `sign` is not +1 or -1, or if `u` contains non-finite values.
+    """
+    u = np.asarray(u, dtype=float)
+    _require_finite("u", u)
+    s = _normalize_sign(sign)
+    return s * (u ** 2)
+
+
+def sqrt_derivatives_to_x_at_zero(
+    order: int,
+    x0: float,
+    g2: Optional[np.ndarray] = None,
+    g4: Optional[np.ndarray] = None,
+) -> np.ndarray:
+    """Pull back derivatives at value x0=0 from u-space (sqrt-domain) to physical x.
+
+    This method maps derivatives computed in the internal sqrt-domain coordinate u
+    back to physical coordinates x at the expansion point x0=0 using the chain rule.
+
+    Args:
+        order: Derivative order to return (1 or 2).
+        x0: Expansion point in physical coordinates (finite). May be +0.0 or -0.0 to
+            select the domain side at the boundary. The domain sign is inferred
+            solely from x0 (including the sign of zero).
+        g2: Second derivative of g with respect to u at u=0; required for order=1.
+        g4: Fourth derivative of g with respect to u at u=0; required for order=2.
+
+    Returns:
+        The derivative(s) in physical coordinates at x0=0.
+
+    Raises:
+        ValueError: If required inputs (g2/g4) are missing for the requested order.
+        NotImplementedError: If `order` not in {1, 2}.
+    """
+    s = _sgn_from_x0(x0)
+    if order == 1:
+        if g2 is None:
+            raise ValueError("order=1 conversion requires g2 (g'' at u=0).")
+        return np.asarray(g2, dtype=float) / (2.0 * s)
+    if order == 2:
+        if g4 is None:
+            raise ValueError("order=2 conversion requires g4 (g'''' at u=0).")
+        return np.asarray(g4, dtype=float) / (12.0 * s * s)
+    raise NotImplementedError("sqrt_derivatives_to_x_at_zero supports orders 1 and 2.")
+
+
+def _normalize_sign(s: float) -> float:
+    """Validate and normalize a sign value to exactly +1.0 or -1.0.
+
+    Args:
+        s: Input sign value (must be approximately ±1).
+
+    Returns:
+        +1.0 or -1.0.
+
+    Raises:
+        ValueError: If s is not finite or not approximately ±1.
+    """
+    if not np.isfinite(s):
+        raise ValueError("sign must be finite.")
+    if np.isclose(abs(s), 1.0, rtol=0.0, atol=1e-12):
+        return 1.0 if s > 0.0 else -1.0
+    raise ValueError("sign must be +1 or -1.")
+
+
+def _require_finite(name: str, arr: np.ndarray) -> None:
+    """Raises a ``ValueError`` if an array contains any non-finite values.
+
+    Args:
+        name: Name of the array (for error message).
+        arr: Array to check.
+
+    Raises:
+        ValueError: If arr contains any non-finite values.
+    """
+    if not np.all(np.isfinite(arr)):
+        raise ValueError(f"{name} must be finite.")
+
+
+def _sgn_from_x0(x0: float) -> float:
+    """Returns the sign of x0, disambiguating zero using np.signbit."""
+    if not np.isfinite(x0):
+        raise ValueError("x0 must be finite.")
+    if x0 > 0.0:
+        return 1.0
+    if x0 < 0.0:
+        return -1.0
+        # x0 == 0.0: disambiguate +0.0 vs -0.0
+    return -1.0 if np.signbit(x0) else 1.0
+

derivkit/derivatives/autodiff/__init__.py

@@ -0,0 +1 @@
+"""JAX autodiff backend for DerivativeKit."""

derivkit/derivatives/autodiff/jax_autodiff.py

@@ -0,0 +1,95 @@
+r"""JAX-based autodiff backend for DerivativeKit.
+
+This backend is intentionally minimal: it only supports scalar derivatives
+$f: R \mapsto R$ via JAX autodiff, and must be registered explicitly as explained in the example.
+
+Example:
+--------
+Basic usage (opt-in registration):
+
+    >>> from derivkit.derivative_kit import DerivativeKit  # doctest: +SKIP
+    >>> from derivkit.derivatives.autodiff.jax_autodiff import register_jax_autodiff_backend  # doctest: +SKIP
+    >>> register_jax_autodiff_backend()  # doctest: +SKIP
+    >>>
+    >>> def func(x):  # doctest: +SKIP
+    ...     import jax.numpy as jnp
+    ...     return jnp.sin(x) + 0.5 * x**2
+    ...
+    >>> dk = DerivativeKit(func, 1.0)  # doctest: +SKIP
+    >>> dk.differentiate(method="autodiff", order=1)  # doctest: +SKIP
+    >>> dk.differentiate(method="autodiff", order=2)  # doctest: +SKIP
+
+Notes:
+------
+- This backend is scalar-only. For gradients/Jacobians/Hessians of functions
+  with vector inputs/outputs, use the standalone helpers in
+  ``derivkit.autodiff.jax_core`` (e.g. ``autodiff_gradient``).
+- To enable this backend, install the JAX extra: ``pip install "derivkit[jax]"``.
+"""
+
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from derivkit.derivative_kit import register_method
+from derivkit.derivatives.autodiff.jax_core import autodiff_derivative
+from derivkit.derivatives.autodiff.jax_utils import require_jax
+
+__all__ = [
+    "AutodiffDerivative",
+    "register_jax_autodiff_backend",
+]
+
+
+class AutodiffDerivative:
+    """DerivativeKit engine for JAX-based autodiff.
+
+    Supports scalar functions f: R -> R with JAX-differentiable bodies.
+    """
+
+    def __init__(self, function: Callable[[float], Any], x0: float):
+        """Initializes the JAX autodiff derivative engine."""
+        self.function = function
+        self.x0 = float(x0)
+
+    def differentiate(self, *, order: int = 1, **_: Any) -> float:
+        """Computes the k-th derivative via JAX autodiff.
+
+        Args:
+            order: Derivative order (>=1).
+
+        Returns:
+            Derivative value as a float.
+        """
+        return autodiff_derivative(self.function, self.x0, order=order)
+
+
+def register_jax_autodiff_backend(
+    *,
+    name: str = "autodiff",
+    aliases: tuple[str, ...] = ("jax", "jax-autodiff", "jax-diff", "jd"),
+) -> None:
+    """Registers the experimental JAX autodiff backend with DerivativeKit.
+
+    After calling this, you can use:
+
+        DerivativeKit(f, x0).differentiate(method=name, order=...)
+
+    Args:
+        name: Name of the method to register.
+        aliases: Alternative names for the method.
+
+    Returns:
+        None
+
+    Raises:
+        AutodiffUnavailable: If JAX is not available.
+    """
+    require_jax()
+
+    register_method(
+        name=name,
+        cls=AutodiffDerivative,
+        aliases=aliases,
+    )

derivkit/derivatives/autodiff/jax_core.py

@@ -0,0 +1,217 @@
+r"""JAX-based autodiff helpers for DerivKit.
+
+This module does not register any DerivKit backend by default.
+
+Use these functions directly, or see
+:func:`derivkit.autodiff.jax_autodiff.register_jax_autodiff_backend`
+for an opt-in integration.
+
+Use only with JAX-differentiable functions. For arbitrary models, prefer
+the "adaptive" or "finite" methods.
+
+Shape conventions (aligned with :mod:`derivKit.calculus` builders):
+
+- ``autodiff_derivative``:
+  :math:`f:\\mathbb{R}\\mapsto\\mathbb{R}` → returns ``float`` (scalar)
+
+- ``autodiff_gradient``:
+  :math:`f:\\mathbb{R}^n\\mapsto\\mathbb{R}` → returns array of shape ``(n,)``
+
+- ``autodiff_jacobian``:
+  :math:`f:\\mathbb{R}^n\\mapsto\\mathbb{R}^m` (or tensor output) → returns array of
+  shape ``(m, n)``, where ``m = \\prod(\text{out\\_shape})``
+
+- ``autodiff_hessian``:
+  :math:`f:\\mathbb{R}^n\\mapsto\\mathbb{R}` → returns array of shape ``(n, n)``
+"""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable
+
+import numpy as np
+
+from derivkit.derivatives.autodiff.jax_utils import (
+    AutodiffUnavailable,
+    apply_array_nd,
+    apply_scalar_1d,
+    apply_scalar_nd,
+    jax,
+    jnp,
+    require_jax,
+)
+
+__all__ = [
+    "autodiff_derivative",
+    "autodiff_gradient",
+    "autodiff_jacobian",
+    "autodiff_hessian",
+]
+
+
+def autodiff_derivative(func: Callable, x0: float, order: int = 1) -> float:
+    """Calculates the k-th derivative of a function f: R -> R via JAX autodiff.
+
+    Args:
+        func: Callable mapping float -> scalar.
+        x0: Point at which to evaluate the derivative.
+        order: Derivative order (>=1); uses repeated grad for higher orders.
+
+    Returns:
+        Derivative value as a float.
+
+    Raises:
+        AutodiffUnavailable: If JAX is not available or function is not differentiable.
+        ValueError: If order < 1.
+        TypeError: If func(x) is not scalar-valued.
+    """
+    require_jax()
+
+    if order < 1:
+        raise ValueError("autodiff_derivative: order must be >= 1.")
+
+    f_jax = partial(apply_scalar_1d, func, "autodiff_derivative")
+
+    g = f_jax
+    for _ in range(order):
+        g = jax.grad(g)
+
+    try:
+        val = g(x0)
+    except (TypeError, ValueError) as exc:
+        raise AutodiffUnavailable(
+            "autodiff_derivative: function is not JAX-differentiable at x0. "
+            "Use JAX primitives / jax.numpy or fall back to 'adaptive'/'finite'."
+        ) from exc
+
+    return float(val)
+
+
+def autodiff_gradient(func: Callable, x0) -> np.ndarray:
+    """Computes the gradient of a scalar-valued function f: R^n -> R via JAX autodiff.
+
+    Args:
+        func: Function to be differentiated.
+        x0: Point at which to evaluate the gradient.
+
+    Returns:
+        A gradient vector as a 1D numpy.ndarray with shape (n,).
+
+    Raises:
+        AutodiffUnavailable: If JAX is not available or function is not differentiable.
+        TypeError: If func(theta) is not scalar-valued.
+    """
+    require_jax()
+
+    x0_arr = np.asarray(x0, float).ravel()
+
+    f_jax = partial(apply_scalar_nd, func, "autodiff_gradient")
+    grad_f = jax.grad(f_jax)
+
+    try:
+        g = grad_f(jnp.asarray(x0_arr))
+    except (TypeError, ValueError) as exc:
+        raise AutodiffUnavailable(
+            "autodiff_gradient: function is not JAX-differentiable."
+        ) from exc
+
+    return np.asarray(g, dtype=float).reshape(-1)
+
+
+def autodiff_jacobian(
+    func: Callable,
+    x0,
+    *,
+    mode: str | None = None,
+) -> np.ndarray:
+    """Calculates the Jacobian of a vector-valued function via JAX autodiff.
+
+    Output convention matches DerivKit Jacobian builders: we flatten the function
+    output to length m = prod(out_shape), and return a 2D Jacobian of shape (m, n),
+    with n = input dimension.
+
+    Args:
+        func: Function to be differentiated.
+        x0: Point at which to evaluate the Jacobian; array-like, shape (n,).
+        mode: Differentiation mode; None (auto), 'fwd', or 'rev'.
+            If None, chooses 'rev' if m <= n, else 'fwd'. For more details about
+            modes, see JAX documentation for `jax.jacrev` and `jax.jacfwd`.
+
+    Returns:
+        A Jacobian matrix as a 2D numpy.ndarray with shape (m, n).
+
+    Raises:
+        AutodiffUnavailable: If JAX is not available or function is not differentiable.
+        ValueError: If mode is invalid.
+        TypeError: If func(theta) is scalar-valued.
+    """
+    require_jax()
+
+    x0_arr = np.asarray(x0, float).ravel()
+    x0_jax = jnp.asarray(x0_arr)
+
+    f_jax = partial(apply_array_nd, func, "autodiff_jacobian")
+
+    try:
+        y0 = f_jax(x0_jax)
+    except (TypeError, ValueError) as exc:
+        raise AutodiffUnavailable(
+            "autodiff_jacobian: function is not JAX-differentiable at x0."
+        ) from exc
+
+    in_dim = x0_arr.size
+    out_dim = int(np.prod(y0.shape))
+
+    if mode is None:
+        use_rev = out_dim <= in_dim
+    elif mode == "rev":
+        use_rev = True
+    elif mode == "fwd":
+        use_rev = False
+    else:
+        raise ValueError("autodiff_jacobian: mode must be None, 'fwd', or 'rev'.")
+
+    jac_fun = jax.jacrev if use_rev else jax.jacfwd
+
+    try:
+        jac = jac_fun(f_jax)(x0_jax)
+    except (TypeError, ValueError) as exc:
+        raise AutodiffUnavailable(
+            "autodiff_jacobian: failed to trace function with JAX."
+        ) from exc
+
+    jac_np = np.asarray(jac, dtype=float)
+    return jac_np.reshape(out_dim, in_dim)
+
+
+def autodiff_hessian(func: Callable, x0) -> np.ndarray:
+    """Calculates the full Hessian of a scalar-valued function.
+
+    Args:
+        func: A function to be differentiated.
+        x0: Point at which to evaluate the Hessian; array-like, shape (n,) with
+            n = input dimension.
+
+    Returns:
+        A Hessian matrix as a 2D numpy.ndarray with shape (n, n).
+
+    Raises:
+        AutodiffUnavailable: If JAX is not available or function is not differentiable.
+        TypeError: If func(theta) is not scalar-valued.
+    """
+    require_jax()
+
+    x0_arr = np.asarray(x0, float).ravel()
+    x0_jax = jnp.asarray(x0_arr)
+
+    f_jax = partial(apply_scalar_nd, func, "autodiff_hessian")
+
+    try:
+        hess = jax.hessian(f_jax)(x0_jax)
+    except (TypeError, ValueError) as exc:
+        raise AutodiffUnavailable(
+            "autodiff_hessian: function is not JAX-differentiable."
+        ) from exc
+
+    return np.asarray(hess, dtype=float)

derivkit/derivatives/autodiff/jax_utils.py

@@ -0,0 +1,146 @@
+"""Utilities for JAX-based autodiff in DerivKit."""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+try:
+    import jax
+    import jax.numpy as jnp
+except ImportError:
+    jax = None
+    jnp = None
+    _HAS_JAX = False
+else:
+    _HAS_JAX = True
+
+has_jax: bool = _HAS_JAX
+
+__all__ = [
+    "AutodiffUnavailable",
+    "require_jax",
+    "to_jax_scalar",
+    "to_jax_array",
+    "apply_scalar_1d",
+    "apply_scalar_nd",
+    "apply_array_nd",
+]
+
+
+class AutodiffUnavailable(RuntimeError):
+    """Raises when JAX-based autodiff is unavailable."""
+
+
+def require_jax() -> None:
+    """Raises if JAX is not available.
+
+    Args:
+        None.
+
+    Returns:
+        None.
+
+    Raises:
+        AutodiffUnavailable: If JAX is not installed.
+    """
+    if not _HAS_JAX:
+        raise AutodiffUnavailable(
+            "JAX autodiff requires `jax` + `jaxlib`.\n"
+            'Install with `pip install "derivkit[jax]"` '
+            "(or follow JAX's official install instructions for GPU)."
+        )
+
+
+def to_jax_scalar(y: Any, *, where: str) -> jnp.ndarray:
+    """Ensures that output is scalar and returns as JAX array.
+
+    Args:
+        y: Output to check.
+        where: Context string for error messages.
+
+    Returns:
+        Scalar (0-d) JAX array with shape ().
+
+    Raises:
+        TypeError: If output is not scalar.
+    """
+    arr = jnp.asarray(y)
+    if arr.ndim != 0:
+        raise TypeError(f"{where}: expected scalar output; got shape {tuple(arr.shape)}.")
+    return arr
+
+
+def to_jax_array(y: Any, *, where: str) -> jnp.ndarray:
+    """Ensures that output is array-like (not scalar) and returns as JAX array.
+
+    Args:
+        y: Output to check.
+        where: Context string for error messages.
+
+    Returns:
+        Non-scalar JAX array with shape (m,) or higher.
+
+    Raises:
+        TypeError: If output is scalar or cannot be converted to JAX array.
+    """
+    try:
+        arr = jnp.asarray(y)
+    except TypeError as exc:
+        raise TypeError(f"{where}: output could not be converted to a JAX array.") from exc
+    if arr.ndim == 0:
+        raise TypeError(f"{where}: output is scalar; use autodiff_derivative/gradient instead.")
+    return arr
+
+
+def apply_scalar_1d(
+    func: Callable[[float], Any],
+    where: str,
+    x: jnp.ndarray,
+) -> jnp.ndarray:
+    """Takes an input function and maps it over a 1D array with scalar output enforcement.
+
+    Args:
+        func: Function to apply.
+        where: Context string for error messages.
+        x: 1D JAX array of inputs.
+
+    Returns:
+        JAX array of scalar outputs.
+    """
+    return to_jax_scalar(func(x), where=where)
+
+
+def apply_scalar_nd(
+    func: Callable,
+    where: str,
+    theta: jnp.ndarray,
+) -> jnp.ndarray:
+    """Takes an input function and maps it over an ND array with scalar output enforcement.
+
+    Args:
+        func: Function to apply.
+        where: Context string for error messages.
+        theta: ND JAX array of inputs.
+
+    Returns:
+        JAX array of scalar outputs.
+    """
+    return to_jax_scalar(func(theta), where=where)
+
+
+def apply_array_nd(
+    func: Callable,
+    where: str,
+    theta: jnp.ndarray,
+) -> jnp.ndarray:
+    """Takes an input function and maps it over an ND array with array output enforcement.
+
+    Args:
+        func: Function to apply.
+        where: Context string for error messages.
+        theta: ND JAX array of inputs.
+
+    Returns:
+        JAX array of array outputs.
+    """
+    return to_jax_array(func(theta), where=where)