bs-python-utils 0.0.1__py3-none-any.whl

bs_python_utils/bsutils.py

@@ -0,0 +1,363 @@
+"""
+Contains various utilities programs.
+"""
+import sys
+import traceback
+from functools import wraps
+from io import TextIOBase
+from math import exp, factorial, log, sqrt
+from pathlib import Path
+from typing import Any, Callable, Iterable, cast
+
+TwoFloats = tuple[float, float]
+ThreeFloats = tuple[float, float, float]
+
+
+def printargs(func: Callable) -> Callable:
+    """
+    Decorator that reports the arguments of the function
+
+    Args:
+      func: the decorated function
+    """
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        print(
+            f"Function {func.__name__} called with args = {args} and kwargs = {kwargs}"
+        )
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def bs_name_func(back: int = 2) -> str:
+    """
+    get the name of the current function, or further back in stack
+
+    Args:
+        back: 2 is current function, 3 the function that called it etc
+
+    Returns:
+        the name of the function
+    """
+    stack = traceback.extract_stack()
+    *_, func_name, _ = stack[-back]
+    return cast(str, func_name)
+
+
+def bs_error_abort(msg: str = "error, aborting") -> None:
+    """
+    report error and abort
+
+    Args:
+        msg: a message
+
+    Returns:
+        prints the message and exits with code 1
+    """
+    print_stars(f"{bs_name_func(3)}: {msg}")
+    sys.exit(1)
+
+
+def bs_switch(
+    match: str, dico: dict, strict: bool = True, default: Any = "no match"
+) -> Any:
+    """
+    a switch statement that allows for partial matches if strict is False
+
+    Args:
+        match: what we look for in the keys
+        dico: a dictionary with string keys
+        strict: if `False`, we accept a partial match
+        default: what we return if no match is found
+
+    Returns:
+        the value for the match, or `default`
+
+    Example:
+        calc_dict = {
+        "plus": lambda x, y: x + y,
+        "minus": lambda x, y: x - y
+        }
+
+        plus = bs_switch('plus', calc_dict, default="unintended function")
+        minus = bs_switch('min', calc_dict, strict=False, default="unintended function")
+
+        plus(6, 4)
+        >> 10
+        minus(6, 4)
+        >> 2
+
+        bs_switch('plu', calc_dict)
+        >> "no match""
+
+    """
+    if strict:
+        for key in dico:
+            if match == key:
+                return dico.get(key)
+    else:
+        for key in dico:
+            if match in key:
+                return dico.get(key)
+    return default
+
+
+def find_first(iterable: Iterable, condition: Callable = lambda x: True) -> Any:
+    """
+    Returns the index of and the first item in the `iterable` that
+    satisfies the `condition`.
+
+    Args:
+        iterable: where to look
+        condition: must return a boolean
+
+    Returns:
+        If the condition is not given, returns 0 and the first item of
+        the iterable.
+
+        Raises `StopIteration` if no item satisfyng the condition is found.
+
+    Example:
+        >>> find_first( (1,2,3), condition=lambda x: x % 2 == 0)
+        (1, 2)
+        >>> find_first(range(3, 100))
+        (0, 3)
+        >>> find_first( () )
+        Traceback (most recent call last):
+        ...
+        StopIteration
+
+
+    """
+    return next((i, x) for i, x in enumerate(iterable) if condition(x))
+
+
+def print_stars(title: str = None, n: int = 70) -> None:
+    """
+    prints a title within stars
+
+    Args:
+        title:  title
+        n: number of stars on line
+
+    Returns:
+        prints a starred line, or two around the title
+    """
+    line_stars = "*" * n
+    print()
+    print(line_stars)
+    if title:
+        print(title.center(n))
+        print(line_stars)
+    print()
+
+
+def file_print_stars(file_handle: TextIOBase, title: str = None, n: int = 70) -> None:
+    """
+    prints to a file a title within stars
+
+    Args:
+        file_handle: file handle
+        title:  title
+        n:   length of line
+
+    Returns:
+        prints a starred line to the file, or two around the title
+    """
+    line_stars = "*" * n
+    file_handle.write("\n")
+    file_handle.write(line_stars)
+    file_handle.write("\n")
+    if title:
+        file_handle.write(title.center(n))
+        file_handle.write("\n")
+        file_handle.write(line_stars)
+        file_handle.write("\n")
+    file_handle.write("\n")
+
+
+def mkdir_if_needed(p: Path | str) -> Path:
+    """
+    create the directory if it does not exist
+
+    Args:
+        p: a path
+
+    Returns:
+        the directory Path
+
+    """
+    try:
+        q = Path(p)
+    except OSError:
+        bs_error_abort(f"{p} is not a path")
+    if not q.exists():
+        q.mkdir(parents=True)
+    return q
+
+
+def bscomb(n: int, k: int) -> int:
+    """
+    number of combinations of k among n `{n \\choose k}`
+
+    Args:
+        n:
+        k: should be smaller than n
+
+    Returns:
+        `{n \\choose k}`
+    """
+    if not isinstance(n, int):
+        bs_error_abort(f"n should be an integer, not {n}")
+    if not isinstance(k, int):
+        bs_error_abort(f"k should be an integer, not {k}")
+    if n < k:
+        bs_error_abort(f"k={k} should not be larger than n={n}")
+    return factorial(n) // (factorial(k) * factorial(n - k))
+
+
+def bslog(
+    x: float, eps: float = 1e-30, deriv: int = 0
+) -> float | TwoFloats | ThreeFloats:
+    """
+    extends the logarithm below `eps` by taking a second-order approximation
+    perhaps with derivatives
+
+    Args:
+        x: argument
+        eps: lower bound
+        deriv: if 1, also return first derivative; if 2, the first two derivatives
+
+    Returns:
+        `\\ln(x)` `C^2`-extended below `eps`, perhaps with derivatives
+    """
+    if deriv not in [0, 1, 2]:
+        bs_error_abort(f"deriv can only be 0, 1, or 2; not {deriv}")
+    if x > eps:
+        logx = log(x)
+        if deriv == 0:
+            return logx
+        dlogx = 1.0 / x
+        if deriv == 1:
+            return logx, dlogx
+        d2logx = -dlogx * dlogx
+        return logx, dlogx, d2logx
+    else:
+        dx = 1.0 - x / eps
+        log_smaller = log(eps) - dx - dx * dx / 2.0
+        if deriv == 0:
+            return log_smaller
+        dlog_smaller = (1.0 + dx) / eps
+        if deriv == 1:
+            return log_smaller, dlog_smaller
+        d2log_smaller = -1.0 / eps / eps
+        return log_smaller, dlog_smaller, d2log_smaller
+
+
+def bsxlogx(
+    x: float, eps: float = 1e-30, deriv: int = 0
+) -> float | TwoFloats | ThreeFloats:
+    """
+    extends `x \\ln(x)` below `eps` by making it go to zero in a `C^2` extension
+    perhaps with derivatives
+
+    Args:
+        x: argument
+        eps: lower bound
+        deriv: if 1, also return first derivative; if 2, the first two derivatives
+
+    Returns:
+        `x \\ln(x)`  `C^2`-extended below `eps`, perhaps with derivatives
+    """
+    if deriv not in [0, 1, 2]:
+        bs_error_abort(f"deriv can only be 0, 1, or 2; not {deriv}")
+    if x > eps:
+        logx = log(x)
+        if deriv == 0:
+            return x * logx
+        if deriv == 1:
+            return x * logx, 1.0 + logx
+        return x * logx, 1.0 + logx, 1.0 / x
+    else:
+        logeps = log(eps)
+        dx = x / eps
+        log_smaller = x * logeps - eps / 2.0 + x * dx / 2.0
+        if deriv == 0:
+            return log_smaller
+        if deriv == 1:
+            return log_smaller, logeps + dx
+        return log_smaller, logeps + dx, 1.0 / eps
+
+
+def _bsexp_extend(x: float, deriv: int, limx: float) -> float | TwoFloats | ThreeFloats:
+    """extends the exponential C^2-wise beyond limx"""
+    elimx = exp(limx)
+    dx = x - limx
+    exp_extend = elimx * (1.0 + dx * (1.0 + 0.5 * dx))
+    if deriv == 0:
+        return exp_extend
+    dexp_extend = elimx * (1.0 + dx)
+    if deriv == 1:
+        return exp_extend, dexp_extend
+    # deriv = 2
+    return exp_extend, dexp_extend, elimx
+
+
+def bsexp(
+    x: float,
+    bigx: float = 50.0,
+    lowx: float = -50.0,
+    deriv: int = 0,
+) -> float | TwoFloats | ThreeFloats:
+    """
+    `C^2`-extends the exponential above `bigx` and below `lowx`
+    perhaps with derivatives
+
+    Args:
+        x: argument
+        bigx: upper bound
+        lowx: lower bound
+        deriv: if 1, also return first derivative; if 2, the first two derivatives
+
+    Returns:
+        the exponential `C^2`-extended above `bigx` and below `lowx`
+        perhaps with derivatives
+    """
+    if deriv not in [0, 1, 2]:
+        bs_error_abort(f"deriv can only be 0, 1, or 2; not {deriv}")
+    if lowx < x < bigx:
+        expx = exp(x)
+        if deriv == 0:
+            return expx
+        if deriv == 1:
+            return expx, expx
+        return expx, expx, expx
+    elif x < lowx:
+        return _bsexp_extend(x, deriv, lowx)
+    else:
+        return _bsexp_extend(x, deriv, bigx)
+
+
+def bs_projection_point(
+    x: float, y: float, a: float, b: float, c: float
+) -> tuple[float, float, float]:
+    """
+    projection of point (x,y) on line ax+by+c=0
+
+    Args:
+        x: y: coordinates
+        a: b: c: line parameters (as in ax+by+c=0)
+
+    Returns:
+        x_proj: y_proj: coordinates of projection point
+        dist: distance of point from line
+    """
+    a2b2 = a * a + b * b
+    denom = sqrt(a2b2)
+    value = a * x + b * y + c
+    x_proj = x - a * value / a2b2
+    y_proj = y - b * value / a2b2
+    dist = abs(value) / denom
+    return x_proj, y_proj, dist

bs_python_utils/distance_covariances.py

@@ -0,0 +1,258 @@
+"""
+distance covariance and partial distance covariance (Szekely and Rizzo)
+evaluation and tests of independence and conditional independence
+"""
+
+from dataclasses import dataclass
+from math import sqrt
+from typing import cast
+
+import numpy as np
+
+from bs_python_utils.bsnputils import check_square, check_vector_or_matrix
+from bs_python_utils.bsutils import bs_error_abort
+
+
+@dataclass
+class DcovResults:
+    dcov: float
+    dcov_stat: float
+    dcor: float
+    X_dd: np.ndarray
+    Y_dd: np.ndarray
+    unbiased: bool
+
+
+@dataclass
+class PdcovResults:
+    pdcov: float
+    pdcov_stat: float
+    pdcor: float
+    X_dd: np.ndarray
+    Y_dd: np.ndarray
+    Z_dd: np.ndarray
+
+
+def _compute_distances(T: np.ndarray) -> np.ndarray:
+    """
+    compute the Euclidian norms (or absolute values)
+    of all row differences `T_k - T_l`
+
+    Args:
+        T: a vector or a matrix
+
+    Returns:
+        the matrix of norms of differences
+    """
+    ndims_T = check_vector_or_matrix(T, "_compute_distances")
+    if ndims_T == 1:
+        return cast(np.ndarray, np.abs(np.subtract.outer(T, T)))
+    else:
+        n, nv = T.shape
+        A = np.zeros((n, n))
+        for iv in range(nv):
+            Tiv = T[:, iv]
+            Aiv = np.subtract.outer(Tiv, Tiv)
+            A += Aiv * Aiv
+        return np.sqrt(A)
+
+
+def _double_decenter(A: np.ndarray, unbiased: bool = False) -> np.ndarray:
+    """
+    does double decentering on a square matrix A
+
+    Args:
+        A: a matrix
+        unbiased: if `True`, we use the Szekely and Rizzo 2014 formula
+
+    Returns:
+        the doubly decentered matrix
+    """
+    n = check_square(A, "_double_decenter")
+    A_1 = np.sum(A, 0)
+    A_2 = np.sum(A, 1).reshape((-1, 1))
+    A_0 = np.sum(A_1)
+    fac2 = (n - 2) if unbiased else n
+    fac1 = (n - 1) if unbiased else n
+    A_dd = A - A_1 / fac2 - A_2 / fac2 + A_0 / (fac1 * fac2)
+    if unbiased:
+        np.fill_diagonal(A_dd, np.zeros(n))
+    return cast(np.ndarray, A_dd)
+
+
+def _dcov_prod(A: np.ndarray, B: np.ndarray, unbiased: bool = False) -> float:
+    n = check_square(A, "_dcov_prod")
+    m = check_square(B, "_dcov_prod")
+    if m == n:
+        fac3 = (n - 3) if unbiased else n
+        return cast(float, np.sum(A * B) / (n * fac3))
+    else:
+        bs_error_abort("A and B should be square matrices of the same size")
+        return 0.0  # for mypy
+
+
+def dcov_dcor(X: np.ndarray, Y: np.ndarray, unbiased: bool = False) -> DcovResults:
+    """
+    evaluate the distance covariance and correlation of `X` and `Y`
+
+    Args:
+        X: `n` observations of a random variable or vector
+        Y: `n` observations of a random variable or vector
+        unbiased: if `True`, we use the Szekely and Rizzo 2014 formula
+
+    Returns:
+        `dCov^2(X,Y)` and `dCor^2(X,Y)`
+    """
+    X_dist = _compute_distances(X)
+    n = X_dist.shape[0]
+    X_dd = _double_decenter(X_dist, unbiased)
+    Y_dist = _compute_distances(Y)
+    Y_dd = _double_decenter(Y_dist, unbiased)
+    dcov2 = _dcov_prod(X_dd, Y_dd, unbiased)
+    dcor2 = dcov2 / sqrt(
+        _dcov_prod(X_dd, X_dd, unbiased) * _dcov_prod(Y_dd, Y_dd, unbiased)
+    )
+    return DcovResults(
+        dcov=dcov2,
+        dcor=dcor2,
+        X_dd=X_dd,
+        Y_dd=Y_dd,
+        unbiased=unbiased,
+        dcov_stat=n * dcov2,
+    )
+
+
+def _dcov_bootstrap(
+    X_dd: np.ndarray,
+    Y_dd: np.ndarray,
+    unbiased: bool = False,
+    ndraws: int = 199,
+) -> np.ndarray:
+    """
+    use bootstrap on the test statistics of independence
+
+    Args:
+        X_dd: the doubly decentered distances for `X`
+        Y_dd: the doubly decentered distances for `Y`
+        unbiased:  if `True`, we use the Szekely and Rizzo 2014 formula
+        ndraws: number of permutations
+
+    Returns:
+        the values of the `ndraws` bootstrapped test stats
+    """
+    n = X_dd.shape[0]
+    dcov_stats_boot = np.zeros(ndraws)
+    for idraw in range(ndraws):
+        draws = np.random.choice(np.arange(n), n)
+        X_ddi = X_dd[draws, :][:, draws]
+        Y_ddi = Y_dd[draws, :][:, draws]
+        if idraw % 50 == 0:
+            print(f"    bootstrap draw {idraw}")
+        dcov_stats_boot[idraw] = _dcov_prod(X_ddi, Y_ddi, unbiased)
+    return cast(np.ndarray, n * dcov_stats_boot)
+
+
+def pvalue_dcov(dcov_results: DcovResults, ndraws: int = 199) -> float:
+    """
+    test of no dependence between `X` and `Y` given `Z`
+
+    Args:
+        dcov_results:  results from `dcov_dcor`
+        ndraws: the number of draws we use
+
+    Returns:
+        the bootstrapped  p-value of the test
+    """
+    X_dd = dcov_results.X_dd
+    Y_dd = dcov_results.Y_dd
+    dcov_stat = dcov_results.dcov_stat
+    unbiased = dcov_results.unbiased
+    dcov_stats_boot = _dcov_bootstrap(X_dd, Y_dd, unbiased, ndraws)
+    sum_small = cast(int, np.sum(dcov_stat < dcov_stats_boot))
+    return (1.0 + sum_small) / (1.0 + ndraws)
+
+
+def pdcov_pdcor(X: np.ndarray, Y: np.ndarray, Z: np.ndarray) -> PdcovResults:
+    """
+    evaluate the partial distance covariance and correlation of `X` and `Y` given `Z`
+
+    Args:
+        X: `n` observations of a random variable or vector
+        Y: `n` observations of a random variable or vector
+        Z: `n` observations of a random variable or vector
+
+    Returns:
+        a `PdcovResults` instance
+    """
+    unbiased = True
+    X_dist = _compute_distances(X)
+    X_dd = _double_decenter(X_dist, unbiased)
+    Y_dist = _compute_distances(Y)
+    Y_dd = _double_decenter(Y_dist, unbiased)
+    Z_dist = _compute_distances(Z)
+    Z_dd = _double_decenter(Z_dist, unbiased)
+    C_XX = _dcov_prod(X_dd, X_dd, unbiased)
+    C_XY = _dcov_prod(X_dd, Y_dd, unbiased)
+    C_YY = _dcov_prod(Y_dd, Y_dd, unbiased)
+    C_XZ = _dcov_prod(X_dd, Z_dd, unbiased)
+    C_YZ = _dcov_prod(Y_dd, Z_dd, unbiased)
+    C_ZZ = _dcov_prod(Z_dd, Z_dd, unbiased)
+    pdcov = C_XY - (C_XZ * C_YZ) / C_ZZ
+    pdcor = pdcov / sqrt((C_XX - C_XZ * C_XZ / C_ZZ) * (C_YY - C_YZ * C_YZ / C_ZZ))
+    n = X.shape[0]
+    return PdcovResults(
+        pdcov=pdcov, pdcor=pdcor, pdcov_stat=n * pdcov, X_dd=X_dd, Y_dd=Y_dd, Z_dd=Z_dd
+    )
+
+
+def _pdcovs_bootstrap(
+    X_dd: np.ndarray, Y_dd: np.ndarray, Z_dd: np.ndarray, ndraws: int = 199
+) -> np.ndarray:
+    """
+    use permutations and recompute the test statistics of independence
+
+    Args:
+        X_dd: the doubly decentered distances for `X`
+        Y_dd: the doubly decentered distances for `Y`
+        Z_dd: the doubly decentered distances for `Y`
+        ndraws: the number of draws we use
+
+    Returns:
+        the `ndraws` values of `pdCov(X,Y ; Z)`
+    """
+    pdcov_stats_boot = np.zeros(ndraws)
+    unbiased = True
+    n = X_dd.shape[0]
+    for idraw in range(ndraws):
+        if idraw % 50 == 0:
+            print(f"pdcov test: bootstrap draw {idraw}")
+        draws = np.random.choice(np.arange(n), n)
+        X_ddi = X_dd[draws, :][:, draws]
+        Y_ddi = Y_dd[draws, :][:, draws]
+        Z_ddi = Z_dd[draws, :][:, draws]
+        C_XY = _dcov_prod(X_ddi, Y_ddi, unbiased)
+        C_XZ = _dcov_prod(X_ddi, Z_ddi, unbiased)
+        C_YZ = _dcov_prod(Y_ddi, Z_ddi, unbiased)
+        C_ZZ = _dcov_prod(Z_ddi, Z_ddi, unbiased)
+        pdcov_stats_boot[idraw] = C_XY - (C_XZ * C_YZ) / C_ZZ
+    return n * pdcov_stats_boot
+
+
+def pvalue_pdcov(pdcov_results: PdcovResults, ndraws: int = 199) -> float:
+    """
+    test of no dependence between `X` and `Y` given `Z`
+
+    Args:
+        pdcov_results: the results of `pdcov_pdcor`
+        ndraws: the number of draws we use
+
+    Returns:
+        the bootstrapped  p-value of the test
+    """
+    X_dd = pdcov_results.X_dd
+    Y_dd = pdcov_results.Y_dd
+    Z_dd = pdcov_results.Z_dd
+    pdcov_stat = pdcov_results.pdcov_stat
+    pdcov_stats_boot = _pdcovs_bootstrap(X_dd, Y_dd, Z_dd, ndraws)
+    sum_small = cast(int, np.sum(pdcov_stat < pdcov_stats_boot))
+    return (1.0 + sum_small) / (1.0 + ndraws)

bs_python_utils/example_opt.py

@@ -0,0 +1,71 @@
+"""examples of optimization"""
+
+import numpy as np
+
+from bs_python_utils.bsutils import print_stars
+from bs_python_utils.bssputils import describe_array
+from bs_python_utils.bs_opt import (
+    acc_grad_descent,
+    minimize_some_fixed,
+    armijo_alpha,
+    barzilai_borwein_alpha,
+)
+
+
+print_stars("Testing acc_grad_descent")
+
+
+def grad_f(x, p):
+    xp = x - p[0]
+    return 4.0 * xp * xp * xp
+
+
+x_init = np.random.normal(size=10000)
+
+p = 1.0
+x_conv, ret_code = acc_grad_descent(
+    grad_f, x_init, other_params=np.array([p]), tol=1e-12, verbose=False
+)
+
+describe_array(x_conv - p, "x-p should be close to zero")
+
+
+def obj(x, args):
+    res = x - args
+    return np.sum(res * res)
+
+
+def grad_obj(x, args):
+    res = x - args
+    return 2.0 * res
+
+
+n = 5
+x_init = np.full(n, 0.5)
+args = np.arange(n)
+bounds = [(-10.0, 10.0) for _ in range(n)]
+
+fixed_vars = [1, 3]
+fixed_vals = -np.ones(2)
+
+resopt = minimize_some_fixed(
+    obj,
+    grad_obj,
+    x_init,
+    args,
+    fixed_vars=fixed_vars,
+    fixed_vals=fixed_vals,
+    bounds=bounds,
+)
+
+print(resopt)
+
+# test the step routines
+g = grad_obj(x_init, args)
+alpha_a = armijo_alpha(obj, x_init, -g, args)
+print(f"\nArmijo alpha={alpha_a}")
+
+alpha_b, g_b = barzilai_borwein_alpha(grad_obj, x_init, args)
+print(f"\nBarzilai-Borwein alpha={alpha_a}")
+print("g and g_b:")
+print(np.column_stack((g, g_b)))