potentials.py

"""Model test potentials for optimization / geometry tests.

Provides:
- rosenbrock(x, a=1.0, b=100.0) -> float
- rosenbrock_grad(x, a=1.0, b=100.0) -> numpy.ndarray

- muller_brown(xy) -> float   # xy is sequence-like of length 2
- muller_brown_grad(xy) -> numpy.ndarray shape (2,)

Implementations use numpy and are small, pure-Python utilities intended for
unit tests and examples.
"""
from __future__ import annotations

import math
from typing import Callable, Sequence, Tuple, Union

try:
    import numpy as np
except Exception:  # pragma: no cover - if numpy missing tests will fail elsewhere
    np = None


def rosenbrock(x: Sequence[float], a: float = 1.0, b: float = 100.0) -> float:
    """Compute the (generalized) Rosenbrock function.

    For a 2D input [x, y], f = (a - x)**2 + b*(y - x**2)**2.
    For n-dimensional input, use successive pairs: sum_{i=0}^{n-2} [(a - x_i)**2 + b*(x_{i+1} - x_i**2)**2].

    Args:
        x: sequence of floats (length >= 2 recommended)
        a: Rosenbrock parameter (default 1.0)
        b: Rosenbrock parameter (default 100.0)

    Returns:
        float: function value
    """
    if np is not None:
        arr = np.asarray(x, dtype=float)
        if arr.size < 2:
            raise ValueError("rosenbrock requires at least 2 variables")
        xi = arr[:-1]
        xnext = arr[1:]
        return float(np.sum((a - xi) ** 2 + b * (xnext - xi ** 2) ** 2))

    # fallback pure-Python
    if len(x) < 2:
        raise ValueError("rosenbrock requires at least 2 variables")
    total = 0.0
    for i in range(len(x) - 1):
        xi = float(x[i])
        xj = float(x[i + 1])
        total += (a - xi) ** 2 + b * (xj - xi ** 2) ** 2
    return total


def rosenbrock_grad(x: Sequence[float], a: float = 1.0, b: float = 100.0):
    """Gradient of the Rosenbrock function.

    Returns a numpy array of same length as x with partial derivatives.
    """
    if np is None:
        raise RuntimeError("numpy is required for rosenbrock_grad")
    arr = np.asarray(x, dtype=float)
    n = arr.size
    if n < 2:
        raise ValueError("rosenbrock_grad requires at least 2 variables")
    grad = np.zeros_like(arr)
    # contributions from terms involving x_i and neighbors
    for i in range(n - 1):
        xi = arr[i]
        xj = arr[i + 1]
        # derivative wrt xi for term i
        grad[i] += -2.0 * (a - xi) - 4.0 * b * xi * (xj - xi * xi)
        # derivative wrt x_{i+1} for term i
        grad[i + 1] += 2.0 * b * (xj - xi * xi)
    return grad

def muller_brown(xy: Sequence[float]) -> float:
    """Evaluate the Müller–Brown potential at 2D position `xy`.

    Parameters are the classical four-term Müller–Brown surface.

    Args:
        xy: sequence (x, y)

    Returns:
        float potential value
    """
    x, y = float(xy[0]), float(xy[1])
    A = [-200.0, -100.0, -170.0, 15.0]
    a = [-1.0, -1.0, -6.5, 0.7]
    b = [0.0, 0.0, 11.0, 0.6]
    c = [-10.0, -10.0, -6.5, 0.6]
    x0 = [1.0, 0.0, -0.5, -1.0]
    y0 = [0.0, 0.5, 1.5, 1.0]
    total = 0.0
    for Ai, ai, bi, ci, xi0, yi0 in zip(A, a, b, c, x0, y0):
        dx = x - xi0
        dy = y - yi0
        total += Ai*math.exp(ai*dx*dx + bi * dx * dy + ci * dy * dy)
    return total

def muller_brown_grad(xy: Sequence[float]) -> 'Tuple[float, float] | np.ndarray':
    """Gradient of the Müller–Brown potential at 2D position `xy`.

    Returns (gx, gy) as a tuple of floats or a numpy array when numpy is available.
    """
    x, y = float(xy[0]), float(xy[1])
    # same parameterization as muller_brown
    A = [-200.0, -100.0, -170.0, 15.0]
    a = [-1.0, -1.0, -6.5, 0.7]
    b = [0.0, 0.0, 11.0, 0.6]
    c = [-10.0, -10.0, -6.5, 0.6]
    x0 = [1.0, 0.0, -0.5, -1.0]
    y0 = [0.0, 0.5, 1.5, 1.0]

    if np is not None:
        A_arr = np.asarray(A, dtype=float)
        a_arr = np.asarray(a, dtype=float)
        b_arr = np.asarray(b, dtype=float)
        c_arr = np.asarray(c, dtype=float)
        x0_arr = np.asarray(x0, dtype=float)
        y0_arr = np.asarray(y0, dtype=float)
        dx = x - x0_arr
        dy = y - y0_arr
        exp_term = np.exp(a_arr * dx * dx + b_arr * dx * dy + c_arr * dy * dy)
        gx = float(np.sum(A_arr * exp_term * (2.0 * a_arr * dx + b_arr * dy)))
        gy = float(np.sum(A_arr * exp_term * (b_arr * dx + 2.0 * c_arr * dy)))
        return np.array([gx, gy], dtype=float)

    # pure Python fallback
    gx = 0.0
    gy = 0.0
    for Ai, ai, bi, ci, xi0, yi0 in zip(A, a, b, c, x0, y0):
        dx = x - xi0
        dy = y - yi0
        ex = math.exp(ai * dx * dx + bi * dx * dy + ci * dy * dy)
        gx += Ai * ex * (2.0 * ai * dx + bi * dy)
        gy += Ai * ex * (bi * dx + 2.0 * ci * dy)
    return (gx, gy)

def visualize_potential(potential: Union[str, Callable[[Sequence[float]], float]], *, xlim=None, ylim=None, vmin=None, vmax=None, nx=400, ny=400, ax=None):
    """Visualize an analytic potential energy surface.

    Parameters:
        potential: either the string name of a supported potential ('muller-brown', 'rosenbrock')
                   or a callable that accepts a sequence-like (x,y) and returns a float.
        xlim: tuple (xmin, xmax). If None, sensible defaults per potential are used.
        ylim: tuple (ymin, ymax). If None, sensible defaults per potential are used.
        vmin, vmax: color scale limits. If None, sensible defaults per potential are used.
        nx, ny: grid resolution (ints)
        ax: optional matplotlib Axes object to draw on; if None, a new figure is created

    Returns:
        (fig, ax): the matplotlib Figure and Axes objects so callers can save or show as they wish.

    Notes:
    - This function does not call plt.show() to remain friendly to headless/test environments.
    """
    # local imports so the module can be imported without matplotlib
    try:
        import matplotlib.pyplot as plt
    except Exception as exc:
        raise RuntimeError("matplotlib is required to visualize potentials") from exc

    if np is None:
        raise RuntimeError("numpy is required to visualize potentials")

    # choose potential function and defaults
    if isinstance(potential, str):
        key = potential.lower()
        if key in ('muller', 'muller-brown', 'muller_brown'):
            potential_func = muller_brown
            default_xlim = (-1.5, 1.5)
            default_ylim = (-0.5, 2.5)
            default_vmin, default_vmax = -200, 250
            contour_levels = np.linspace(default_vmin, default_vmin + 400, 40)
        elif key in ('rosen', 'rosenbrock'):
            potential_func = rosenbrock
            default_xlim = (-1.5, 1.5)
            default_ylim = (-0.5, 2.5)
            default_vmin, default_vmax = 0, 100
            contour_levels = np.linspace(default_vmin, default_vmin + 2, 10)
        else:
            raise ValueError(f"Unknown potential name: {potential}")
    elif callable(potential):
        potential_func = potential
        default_xlim = (-1.5, 1.5)
        default_ylim = (-0.5, 2.5)
        default_vmin, default_vmax = -200, 250
    else:
        raise TypeError("'potential' must be a string name or a callable")

    # apply defaults
    xlim = tuple(xlim) if xlim is not None else default_xlim
    ylim = tuple(ylim) if ylim is not None else default_ylim
    vmin = vmin if vmin is not None else default_vmin
    vmax = vmax if vmax is not None else default_vmax

    # grid
    xs = np.linspace(xlim[0], xlim[1], nx)
    ys = np.linspace(ylim[0], ylim[1], ny)
    X, Y = np.meshgrid(xs, ys)

    # vectorize potential evaluation (safe wrapper to ensure float output)
    vec = np.vectorize(lambda xi, yi: float(potential_func([xi, yi])))
    Z = vec(X, Y)

    created_fig = False
    if ax is None:
        fig, ax = plt.subplots(figsize=(8, 6))
        created_fig = True
    else:
        fig = ax.figure

    levels = np.linspace(vmin, vmax, 100)
    cp = ax.contourf(X, Y, Z, levels=levels, cmap='autumn')
    # add contour lines if we have sensible levels
    if 'contour_levels' in locals():
        ax.contour(X, Y, Z, levels=contour_levels, colors='black', linewidths=0.5, alpha=0.5)
    # only add colorbar if we created the figure here
    # when we created the figure, add a colorbar; otherwise the caller may manage colorbar
    if created_fig:
        fig.colorbar(cp, ax=ax)
    ax.set_title('Potential Energy Surface')
    ax.set_xlabel('X')
    ax.set_ylabel('Y')

    return fig, ax