Source code for jaxmat.tensors.utils

import jax.numpy as jnp
import optax




def safe_fun(fun, x, norm=None, eps=1e-16):
    """
    Safely applies a function to an input, avoiding numerical issues near zero.

    This function applies ``fun(x)`` only when the norm of ``x`` exceeds a small
    tolerance ``eps``. Otherwise, it returns zero. This is useful for ensuring
    numerical stability in cases where evaluating ``fun`` at or near zero could
    result in undefined or unstable behavior (e.g., division by zero).

    Parameters
    ----------
    fun : Callable
        The function to apply safely.
    x : array-like
        Input array or tensor.
    norm : Callable, optional
        A norm or magnitude function used to test whether ``x`` is sufficiently
        large. Defaults to the identity function.
    eps : float, optional
        Small threshold to determine whether ``x`` is treated as nonzero.
        Defaults to 1e-16.

    Returns
    -------
    array-like
        ``fun(x)`` if ``norm(x) > eps``, otherwise ``0`` (of the same shape as ``x``).
    """
    if norm is None:
        norm = lambda x: x
    nonzero_x = jnp.where(norm(x) > eps, x, 0 * x)
    return jnp.where(norm(x) > eps, fun(nonzero_x), 0)





def safe_sqrt(x, eps=1e-16):
    """
    Computes a numerically safe square root.

    Ensures the argument to the square root is greater than `eps`
    to avoid taking the square root of zero or negative values,
    which could cause instability or NaNs.

    Parameters
    ----------
    x : array-like
        Input array or tensor.
    eps : float, optional
        Minimum threshold for `x` before taking the square root. Defaults to 1e-16.

    Returns
    --------
    array-like
        The square root of `x` for `x > eps`, otherwise `eps`.
    """
    nonzero_x = jnp.where(x > eps, x, eps)
    return jnp.where(x > eps, jnp.sqrt(nonzero_x), eps)





def safe_norm(x, eps=1e-16, **kwargs):
    """
    Wrapper around ``optax.safe_norm`` that computes a numerically stable norm.

    This function prevents numerical instability when computing vector norms
    for small magnitudes by internally applying a stability threshold.

    Parameters
    ----------
    x : array-like
        Input vector or tensor.
    eps : float, optional
        Small constant added for numerical stability. Defaults to ``1e-16``.
    **kwargs:
        Additional arguments passed to ``optax.safe_norm``.

    Returns
    -------
    array-like
        The numerically stable norm of ``x``.
    """
    return optax.safe_norm(x, eps, **kwargs)





def FischerBurmeister(x, y):
    r"""
    Computes the scalar Fischer-Burmeister function.

    The Fischer-Burmeister function is defined as:
    $$\Phi(x, y) = x + y - \sqrt{x^2 + y^2}$$

    and is commonly used in complementarity problem formulations to provide
    a smooth reformulation of the complementarity conditions
    $$x \geq 0, y \geq 0, xy = 0$$.
    """
    return x + y - safe_sqrt(x**2 + y**2)