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

bs_python_utils/bs_opt.py

@@ -0,0 +1,518 @@
+""" interface to scipy.optimize """
+from dataclasses import dataclass
+from math import sqrt
+from typing import Any, Callable, Iterable, Optional, Union, cast
+
+import numpy as np
+import scipy.linalg as spla
+import scipy.optimize as spopt
+
+from bs_python_utils.bsnputils import TwoArrays, check_vector, npmaxabs
+from bs_python_utils.bssputils import describe_array
+from bs_python_utils.bsutils import bs_error_abort, print_stars
+from bs_python_utils.Timer import timeit
+
+ScalarFunctionAndGradient = Callable[
+    [np.ndarray, Iterable, Optional[bool]], Union[float, tuple[float, np.ndarray]]
+]
+"""Type of f(v, args, gr) that returns a scalar value and also a gradient if gr is True"""
+
+
+ProximalFunction = Callable[[np.ndarray, float, Iterable], np.ndarray]
+"""Type of h(x, t, pars) that returns a scalar value"""
+
+
+@dataclass
+class OptimizeParams:
+    """
+    used for optimization;
+    combines values, bounds and initial values for a parameter vector
+    """
+
+    params_values: np.ndarray | None
+    params_bounds: list[tuple] | None
+    params_init: np.ndarray | None
+
+
+def print_optimization_results(
+    resus: spopt.OptimizeResult, title: str = "Minimizing"
+) -> None:
+    """
+    print results from unconstrained optimization
+
+    Args:
+        resus: results from optimization
+        title: a title
+
+    Returns:
+        just prints
+    """
+    print_stars(title)
+    print(resus.message)
+    if resus.success:
+        print(f"Successful! in {resus.nit} iterations")
+        print(f" evaluated {resus.nfev} functions functions and {resus.njev} gradients")
+        print("\nMinimizer and grad_f:")
+        print(np.column_stack((resus.x, resus.jac)))
+        print(f"Minimized value is {resus.fun}")
+    else:
+        print_stars("Minimization failed!")
+    return
+
+
+def print_constrained_optimization_results(
+    resus: spopt.OptimizeResult,
+    title: str = "Minimizing",
+    print_constr: bool = False,
+    print_multipliers: bool = False,
+) -> None:
+    """
+    print results from constrained optimization
+
+       Args:
+        resus: results from optimization
+        str title: a title
+        print_constr: if `True`, print the values of the constraints
+        print_multipliers: if `True`, print the values of the multipliers
+
+    Returns:
+        just prints
+    """
+    print_stars(title)
+    print(resus.message)
+    if resus.success:
+        print(f"Successful! in {resus.nit} iterations")
+        print(f" evaluated {resus.nfev} functions and {resus.njev} gradients")
+        print(f"Minimized value is {resus.fun}")
+        print(f"The Lagrangian norm is {resus.optimality}")
+        print(f"The largest constraint violation is {resus.constr_violation}")
+        if print_multipliers:
+            print(f"The multipliers are {resus.v}")
+        if print_constr:
+            print(f"The values of the constraints are {resus.constr}")
+    else:
+        print_stars("Constrained minimization failed!")
+    return
+
+
+def armijo_alpha(
+    f: Callable,
+    x: np.ndarray,
+    d: np.ndarray,
+    args: Iterable,
+    alpha_init: float = 1.0,
+    beta: float = 0.5,
+    max_iter: int = 100,
+    tol: float = 0.0,
+) -> float:
+    """Given a function `f` we are minimizing, computes the step size `alpha`
+    to take in the direction `d` using the Armijo rule
+
+    Args:
+        f: the function
+        x: the current point
+        d: the direction we are taking
+        args: other arguments passed to `f`
+        alpha_init: the initial step size
+        beta: the step size reduction factor
+        max_iter: the maximum number of iterations
+        tol: a tolerance
+
+    Returns:
+        the step size alpha
+    """
+    f0 = f(x, args)
+    alpha = alpha_init
+    for _ in range(max_iter):
+        x1 = x + alpha * d
+        f1 = f(x1, args)
+        if f1 < f0 + tol:
+            return alpha
+        alpha *= beta
+    else:
+        bs_error_abort("Too many iterations")
+    return alpha
+
+
+def barzilai_borwein_alpha(
+    grad_f: Callable, x: np.ndarray, args: Iterable
+) -> tuple[float, np.ndarray]:
+    """Given a function `f` we are minimizing, computes the step size `alpha`
+    to take in the opposite direction of the gradient using the Barzilai-Borwein rule
+
+    Args:
+        grad_f: the gradient of the function
+        x: the current point
+        args: other arguments passed to `f`
+
+    Returns:
+        the step size `alpha` and the gradient `g` at the point `x`
+    """
+    g = grad_f(x, args)
+    alpha = 1.0 / spla.norm(g)
+    x_hat = x - alpha * g
+    g_hat = grad_f(x_hat, args)
+    norm_dg = spla.norm(g - g_hat)
+    norm_dg2 = norm_dg * norm_dg
+    alpha = np.abs(np.dot(x - x_hat, g - g_hat)) / norm_dg2
+    return alpha, g
+
+
+def check_gradient_scalar_function(
+    fg: ScalarFunctionAndGradient,
+    p: np.ndarray,
+    args: Iterable,
+    mode: str = "central",
+    EPS: float = 1e-6,
+) -> TwoArrays:
+    """Checks the gradient of a scalar function
+
+    Args:
+        fg: should return the scalar value, and the gradient if its `gr` argument is `True`
+        p: where we are checking the gradient
+        args: other arguments passed to `fg`
+        mode: "central" or "forward" derivatives
+        EPS: the step for forward or central derivatives
+
+    Returns:
+        the analytic and numeric gradients
+    """
+    f0, f_grad = fg(p, args, gr=True)  # type: ignore
+    f0 = cast(float, f0)
+
+    print_stars("checking the gradient: analytic, numeric")
+
+    g = np.zeros_like(p)
+    if mode == "central":
+        for i, x in enumerate(p):
+            p1 = p.copy()
+            p1[i] = x + EPS
+            f_plus = cast(float, fg(p1, args, gr=False))  # type: ignore
+            p1[i] -= 2.0 * EPS
+            f_minus = cast(float, fg(p1, args, gr=False))  # type: ignore
+            g[i] = (f_plus - f_minus) / (2.0 * EPS)
+            print(f"{i}: {f_grad[i]}, {g[i]}")
+    elif mode == "forward":
+        for i, x in enumerate(p):
+            p1 = p.copy()
+            p1[i] = x + EPS
+            f_plus = cast(float, fg(p1, args, gr=False))  # type: ignore
+            g[i] = (f_plus - f0) / EPS
+            print(f"{i}: {f_grad[i]}, {g[i]}")
+    else:
+        bs_error_abort("mode must be 'central' or 'forward'")
+
+    return f_grad, g
+
+
+@timeit
+def acc_grad_descent(
+    grad_f: Callable,
+    x_init: np.ndarray,
+    other_params: Iterable,
+    prox_h: ProximalFunction | None = None,
+    print_result: bool = False,
+    verbose: bool = False,
+    tol: float = 1e-9,
+    alpha: float = 1.01,
+    beta: float = 0.5,
+    maxiter: int = 10000,
+) -> tuple[np.ndarray, int]:
+    """
+    minimizes `(f+h)` by Accelerated Gradient Descent
+     where `f` is smooth and convex  and `h` is convex.
+
+    By default `h` is zero.
+
+    Args:
+        grad_f: grad_f of `f`; should return an `(n)` array from an `(n)` array \
+        and the `other_ params` object
+        x_init: initial guess, shape `(n)`
+        prox_h: proximal projector of `h`, if any
+        other_params: an iterable with additional parameters
+        verbose: if `True`, print diagnosis
+        tol: convergence criterion on absolute grad_f
+        alpha: ceiling on step multiplier
+        beta: floor on step multiplier
+        maxiter: max number of iterations
+
+    Returns: 
+        the candidate solution, and 1 if converged/0 if not
+    """
+
+    # no proximal projection if no h
+    local_prox_h: ProximalFunction = prox_h if prox_h else lambda x, t, p: x
+
+    x = x_init.copy()
+    y = x_init.copy()
+
+    #  for stepsize we use Barzilai-Borwein
+    t, g = barzilai_borwein_alpha(grad_f, y, other_params)
+
+    grad_err_init = npmaxabs(g)
+
+    if verbose:
+        print(f"agd: grad_err_init={grad_err_init}")
+
+    n_iter = 0
+    theta = 1.0
+
+    while n_iter < maxiter:
+        grad_err = npmaxabs(g)
+        if grad_err < tol:
+            break
+        xi = x
+        yi = y
+        x = y - t * g
+        x = local_prox_h(x, t, other_params)
+
+        theta = 2.0 / (1.0 + sqrt(1.0 + 4.0 / theta / theta))
+
+        if np.dot(y - x, x - xi) > 0:  # wrong direction, we restart
+            x = xi
+            y = x
+            theta = 1.0
+        else:
+            y = x + (1.0 - theta) * (x - xi)
+
+        gi = g
+        g = grad_f(y, other_params)
+        ndy = spla.norm(y - yi)
+        t_hat = 0.5 * ndy * ndy / abs(np.dot(y - yi, gi - g))
+        t = min(alpha * t, max(beta * t, t_hat))
+
+        n_iter += 1
+
+        if verbose:
+            print(f" AGD with grad_err = {grad_err} after {n_iter} iterations")
+
+    x_conv = y
+
+    ret_code = 0 if grad_err < tol else 1
+
+    if verbose or print_result:
+        if ret_code == 0:
+            print_stars(
+                f" AGD converged with grad_err = {grad_err} after {iter} iterations"
+            )
+        else:
+            print_stars(
+                f" Problem in AGD: grad_err = {grad_err} after {iter} iterations"
+            )
+
+    return (x_conv, ret_code)
+
+
+def _fix_some(
+    obj: Callable, grad_obj: Callable, fixed_vars: list[int], fixed_vals: np.ndarray
+) -> tuple[Callable, Callable]:
+    """
+    Takes in a function and its gradient, fixes the variables
+    whose indices are `fixed_vars` to the values in `fixed_vals`,
+    and returns the modified function and its gradient
+
+    Args:
+        obj: the original function
+        grad_obj: its gradient function
+        fixed_vars: a list if the indices of variables whose values are fixed
+        fixed_vals: their fixed values
+
+    Returns:
+        the modified function and its modified gradient function
+    """
+
+    def fixed_obj(t, other_args):
+        t_full = list(t)
+        for i, i_coef in enumerate(fixed_vars):
+            t_full.insert(i_coef, fixed_vals[i])
+        arr_full = np.array(t_full)
+        return obj(arr_full, other_args)
+
+    def fixed_grad_obj(t, other_args):
+        t_full = list(t)
+        for i, i_coef in enumerate(fixed_vars):
+            t_full.insert(i_coef, fixed_vals[i])
+        arr_full = np.array(t_full)
+        grad_full = grad_obj(arr_full, other_args)
+        return np.delete(grad_full, fixed_vars)
+
+    return fixed_obj, fixed_grad_obj
+
+
+def minimize_some_fixed(
+    obj: Callable,
+    grad_obj: Callable,
+    x_init: np.ndarray,
+    args: Iterable,
+    fixed_vars: list[int] | None,
+    fixed_vals: np.ndarray | None,
+    options: dict | None = None,
+    bounds: list[tuple[float, float]] | None = None,
+) -> Any:
+    """
+    minimize a function with some variables fixed, using L-BFGS-B
+
+    Args:
+        obj: the original function
+        grad_obj: its gradient function
+        fixed_vars: a list if the indices of variables whose values are fixed
+        fixed_vals: their fixed values
+        x_init: the initial values of all variables (those on fixed variables are not used)
+        args: other parameters
+        options: any options passed on to scipy.optimize.minimize
+        bounds: the bounds on all variables (those on fixed variables are not used)
+
+    Returns:
+        the result of optimization, on all variables
+    """
+    if fixed_vars is None:
+        resopt = spopt.minimize(
+            obj,
+            x_init,
+            method="L-BFGS-B",
+            args=args,
+            options=options,
+            jac=grad_obj,
+            bounds=bounds,
+        )
+    else:
+        fixed_vars = cast(list, fixed_vars)
+        n_fixed = check_vector(fixed_vals)
+        fixed_vals = cast(np.ndarray, fixed_vals)
+        if len(fixed_vars) != n_fixed:
+            bs_error_abort(
+                f"fixed_vars has {len(fixed_vars)} indices but fixed_vals has"
+                f" {fixed_vals.size} elements."
+            )
+        fixed_obj, fixed_grad_obj = _fix_some(obj, grad_obj, fixed_vars, fixed_vals)
+
+        # drop fixed variables and the corresponding bounds
+        n = len(x_init)
+        not_fixed = np.ones(n, dtype=bool)
+        not_fixed[fixed_vars] = False
+        t_init = x_init[not_fixed]
+        t_bounds = (
+            None if bounds is None else [bounds[i] for i in range(n) if not_fixed[i]]
+        )
+
+        resopt = spopt.minimize(
+            fixed_obj,
+            t_init,
+            method="L-BFGS-B",
+            args=args,
+            options=options,
+            jac=fixed_grad_obj,
+            bounds=t_bounds,
+        )
+
+        # now re-fill the values of the variables
+        t = resopt.x
+        t_full = list(t)
+        for i, i_coef in enumerate(fixed_vars):
+            t_full.insert(i_coef, fixed_vals[i])
+        resopt.x = t_full
+
+        # and re-fill the values of the gradients
+        g = grad_obj(np.array(t_full), args)
+        resopt.jac = g
+
+    return resopt
+
+
+def dfp_update(
+    hess_inv: np.ndarray, gradient_diff: np.ndarray, x_diff: np.ndarray
+) -> np.ndarray:
+    """runs a DFP update for the inverse Hessian
+
+    Args:
+        hess_inv: the current inverse Hessian
+        gradient_diff: the update in the gradient
+        x_diff: the update in x
+
+    Returns:
+        the updated inverse Hessian
+    """
+    xdt = x_diff.T
+    xxp = x_diff @ xdt
+    xpg = xdt @ gradient_diff
+    hdg = hess_inv @ gradient_diff
+    dgp_hdg = gradient_diff.T @ hdg
+    hess_inv_new = hess_inv + xxp / xpg - (hdg @ hdg.T) / dgp_hdg
+    return cast(np.ndarray, hess_inv_new)
+
+
+def bfgs_update(
+    hess_inv: np.ndarray, gradient_diff: np.ndarray, x_diff: np.ndarray
+) -> np.ndarray:
+    """runs a DFP update for the inverse Hessian
+
+    Args:
+        hess_inv: the current inverse Hessian
+        gradient_diff: the update in the gradient
+        x_diff: the update in x
+
+    Returns:
+        the updated inverse Hessian
+    """
+    xdt = x_diff.T
+    xpg = xdt @ gradient_diff
+    hdg = hess_inv @ gradient_diff
+    dgp_hdg = gradient_diff.T @ hdg
+    u = x_diff / xpg - hdg / dgp_hdg
+    hess_inv_new = dfp_update(hess_inv, gradient_diff, x_diff) + dgp_hdg * (u @ u.T)
+    return cast(np.ndarray, hess_inv_new)
+
+
+if __name__ == "__main__":
+    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)))

bs_python_utils/bs_plots.py

@@ -0,0 +1,2 @@
+""" a personal library of plots
+"""

bs_python_utils/bs_seaborn.py

@@ -0,0 +1,174 @@
+""" personal library of Seaborn plots
+"""
+from typing import Callable, cast
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+
+SeabornGraph = tuple[mpl.figure.Figure, mpl.figure.Axes]
+
+
+def bs_sns_get_legend(g: mpl.axes.Axes) -> mpl.legend.Legend:
+    """
+    get the legend object of a Seaborn plot
+
+    Args:
+        g: the plot object
+
+    Returns:
+        leg: the associated Legend object
+    """
+
+    # check axes and find which is have legend
+    axs = g.axes
+    if isinstance(axs, mpl.axes.Axes):  # only one Axes
+        leg = axs.get_legend()
+    else:
+        for ax in axs.flat():
+            leg = ax.get_legend()
+            if leg is not None:
+                break
+    # or legend may be on a figure
+    if leg is None:
+        leg = g._legend
+    return leg
+
+
+def bs_sns_bar_x_byf(
+    df: pd.DataFrame,
+    xstr: str,
+    fstr: str,
+    statistic: Callable = np.mean,
+    label_x: str | None = None,
+    label_f: str | None = None,
+    title: str | None = None,
+) -> SeabornGraph:
+    """make a bar plot of x by f and g
+
+    Args:
+        df: dataframe, should contain columns `xstr` and `fstr`
+        xstr: column name of x
+        fstr: column name of f
+        statistic: statistic to plot (by default, the mean)
+        label_x: label of x
+        label_f: label of f
+        title: title of plot
+
+    Returns:
+        fig: figure
+        ax: axes
+    """
+    fig, ax = plt.subplots()
+    gbar = sns.barplot(
+        x=fstr,
+        y=xstr,
+        data=df,
+        estimator=statistic,
+        errcolor="r",
+        errwidth=0.75,
+        capsize=0.2,
+        ax=ax,
+    )
+    xlab = fstr if label_f is None else label_f
+    ylab = xstr if label_x is None else label_x
+    ax.set_xlabel(xlab)
+    ax.set_ylabel(ylab)
+    if title is not None:
+        ax.set_title(title)
+    return cast(SeabornGraph, gbar)
+
+
+def bs_sns_bar_x_byfg(
+    df: pd.DataFrame,
+    xstr: str,
+    fstr: str,
+    gstr: str,
+    statistic: Callable = np.mean,
+    label_x: str | None = None,
+    label_f: str | None = None,
+    label_g: str | None = None,
+    title: str | None = None,
+) -> SeabornGraph:
+    """make a bar plot of x by f and g
+
+    Args:
+        df: dataframe, should contain columns  `xstr`, `fstr`,  and `gstr`
+        xstr: column name of x
+        fstr: column name of f
+        gstr: column name of g
+        statistic: statistic to plot (by default, the mean)
+        label_x: label of x
+        label_f: label of f
+        label_g: label of g in legend
+        title: title of plot
+
+    Returns:
+        fig: figure
+        ax: axes
+    """
+    fig, ax = plt.subplots()
+    gbar = sns.barplot(
+        x=fstr,
+        y=xstr,
+        data=df,
+        hue=gstr,
+        estimator=statistic,
+        errcolor="r",
+        errwidth=0.75,
+        capsize=0.2,
+        ax=ax,
+    )
+    xlab = fstr if label_f is None else label_f
+    ylab = xstr if label_x is None else label_x
+    ax.set_xlabel(xlab)
+    ax.set_ylabel(ylab)
+    if label_g is not None:
+        gbar_legend = bs_sns_get_legend(gbar)
+        gbar_legend.set_title(label_g)
+    if title is not None:
+        ax.set_title(title)
+    return cast(SeabornGraph, gbar)
+
+
+def bs_sns_density_estimates(
+    df: pd.DataFrame,
+    true_values: np.ndarray,
+    method_string: str | None = "Estimator",
+    coeff_string: str | None = "Parameter",
+    estimate_string: str | None = "Estimate",
+    max_cols: int = 3,
+) -> sns.FacetGrid:
+    """
+    plot of the densities of estimates of several coefficients with several methods,
+    superposed by methods and faceted by coefficients
+
+    Args:
+        df: contains columns `method_string`, `coeff_name`, `estimate_value`
+        true_values: the true values of the coefficients
+        method_string: the name of the column that indicates the method
+        coeff_string: the name of the column that indicates the coefficient
+        estimate_string: the name of the column that gives the value of the estimate
+        max_cols: we wrap after that
+
+    Returns:
+        the `FacetGrid`
+
+    """
+    g = sns.FacetGrid(
+        data=df,
+        sharex=False,
+        sharey=False,
+        hue=method_string,
+        col=coeff_string,
+        col_wrap=max_cols,
+    )
+    g.map(sns.kdeplot, estimate_string)
+    g.set_titles("{col_name}")
+    for true_val, ax in zip(true_values, g.axes.ravel(), strict=True):
+        ax.vlines(true_val, *ax.get_ylim(), color="k", linestyles="dashed")
+    g.add_legend()
+
+    return g