source_code/src/mpm/polynomials.py

"""Orthogonal polynomial utilities for building multiscale systems."""
from __future__ import annotations

from typing import Callable, Tuple

import jax
from jax import jit, vmap, lax
import jax.numpy as jnp
import numpy as np
from numpy.polynomial.legendre import legval as np_legval
from scipy.special import eval_legendre


SystemParams = Tuple[np.ndarray, ...]
EvalFn = Callable[..., np.ndarray]


def get_system_params(
    measure: str,
    n: int,
    m: int | None = None,
    *,
    multiscale: bool = False,
    implicit: bool = False,
    **measure_args,
) -> Tuple[SystemParams, EvalFn, float]:
    if multiscale:
        if m is None:
            raise ValueError("Multiscale systems require `m`")
        return _get_multiscale_system_params(
            measure, n, m, implicit=implicit, **measure_args
        )
    return _transition(measure, n, implicit=implicit, **measure_args)


def get_output_vector(measure: str, n: int, fourier_overlap: float = 0.99) -> np.ndarray:
    if measure == "legt":
        return np.sqrt(2 * np.arange(n) + 1.0) * np.power(-1, np.arange(n))
    if measure == "fout":
        B = np.zeros(n)
        m = 2 * np.pi * np.arange(n // 2)
        B[0::2] = np.sqrt(2.0) * np.cos(m * fourier_overlap)
        B[1::2] = np.sqrt(2.0) * np.sin(m * fourier_overlap)
        B[0] = 1.0
        return B
    raise NotImplementedError(measure)


def _transition(
    measure: str,
    N: int,
    *,
    implicit: bool = False,
    **measure_args,
) -> Tuple[SystemParams, EvalFn, float]:
    if measure == "lagt":
        if implicit:
            raise ValueError("Laguerre systems are not implemented in implicit form")
        b = measure_args.get("beta", 1.0)
        A = np.eye(N) / 2 - np.tril(np.ones((N, N)))
        B = b * np.ones((N, 1))
        eval_func = legval
        tmx_max = 1.0
    elif measure == "fout":
        freqs = 2 * np.arange(N // 2)
        d = np.stack([np.zeros(N // 2), freqs], axis=-1).reshape(-1)[1:]
        A = np.pi * (-np.diag(d, 1) + np.diag(d, -1))
        B = np.zeros(N)
        B[0::2] = np.sqrt(2.0)
        B[0] = 1.0
        A = A - 2 * B[:, None] * B[None, :]
        B = 2 * B[:, None]
        eval_func = fourier_val
        tmx_max = 1.0
    elif measure == "legt":
        Q = np.arange(N, dtype=np.float64)
        R = (2 * Q + 1) ** 0.5
        if implicit:
            a_func = get_implicit_legt_A(N)
            b = R
            params = (a_func, b)
        else:
            j, i = np.meshgrid(Q, Q)
            A = R[:, None] * np.where(i < j, (-1.0) ** (i - j), 1) * R[None, :]
            B = R[:, None]
            A = -A
            params = (A, B.flatten())
        eval_func = legval
        tmx_max = 1.0
    elif measure == "legs":
        if implicit:
            raise ValueError("Scaled Legendre systems do not support implicit mode")
        q = np.arange(N, dtype=np.float64)
        col, row = np.meshgrid(q, q)
        r = 2 * q + 1
        M = -(np.where(row >= col, r, 0) - np.diag(q))
        T = np.sqrt(np.diag(2 * q + 1))
        A = T @ M @ np.linalg.inv(T)
        B = np.diag(T)[:, None]
        eval_func = legsval
        tmx_max = 5.0
        params = (A, B.flatten())
    else:
        raise NotImplementedError(measure)

    if not implicit and measure != "legt":
        params = (A, B.flatten())

    return params, eval_func, tmx_max


def _get_multiscale_system_params(
    measure: str,
    n: int,
    m: int,
    *,
    implicit: bool = False,
    **measure_args,
) -> Tuple[SystemParams, EvalFn, float]:
    params, eval_func, tmx_max = _transition(
        measure, n, implicit=implicit, **measure_args
    )

    if implicit:
        a_func, B = params
        m_func = get_implicit_M(m)
        implicit_A = vmap(a_func, in_axes=1, out_axes=1)
        implicit_M = vmap(m_func, in_axes=1, out_axes=1)
        B = B.reshape(-1, 1)
        B = np.block([B / 2, B / (2 * np.sqrt(3)), np.zeros((n, m - 2))])
        params = (implicit_A, B, implicit_M)
    else:
        A, B = params
        M = make_m_mat(m)
        B = B.reshape(-1, 1)
        B = np.block([B / 2, B / (2 * np.sqrt(3)), np.zeros((n, m - 2))])
        params = (A, B, M)

    def ms_eval_func(tmx, taus, c):
        coeffs = legval(taus, c * np.power(-1, np.arange(c.shape[1]))[None]).T
        return eval_func(tmx, coeffs)

    return params, ms_eval_func, tmx_max * 2


def get_L_vector(m: int, tau: float) -> np.ndarray:
    res = np.zeros(m)
    vec = np.sqrt(2 * np.arange(m) + 1.0)
    for i in range(m):
        oh = np.zeros(m)
        oh[i] = 1.0
        res[i] = np_legval(2 * tau - 1, vec * oh)
    return res


def make_m_mat(m: int) -> np.ndarray:
    out = np.eye(m)
    idx = np.arange(m - 1)
    vec = np.arange(1, m) / np.sqrt(
        np.arange(1, 2 * m - 1, 2) * np.arange(3, 2 * m + 1, 2)
    )
    out[idx, idx + 1] = vec
    out[idx + 1, idx] = vec
    return out / 2


def get_implicit_M(m: int):
    vec = jnp.arange(1, m) / jnp.sqrt(
        jnp.arange(1, 2 * m - 1, 2) * jnp.arange(3, 2 * m + 1, 2)
    )

    @jit
    def matvec(v):
        out = v[:]
        out = out.at[1:].add(vec * v[:-1])
        out = out.at[:-1].add(vec * v[1:])
        return out / 2

    return matvec


def get_implicit_legt_A(n: int):
    if n % 2 != 0:
        raise ValueError("Implicit Legendre systems require an even order")
    if n <= 1:
        raise ValueError("n must be greater than 1")
    sqrt_filter = jnp.sqrt(1.0 + 2.0 * jnp.arange(n))
    alt_filter = (-1) ** jnp.arange(n)

    @jit
    def matvec(v):
        temp = v * sqrt_filter
        alt_cum_sum = jnp.cumsum(temp * alt_filter)
        result = jnp.cumsum(temp) + (alt_cum_sum[-1] - alt_cum_sum) * alt_filter
        return -result * sqrt_filter

    return matvec


def legval(x: np.ndarray, c: np.ndarray) -> np.ndarray:
    """
    Orthonormal Legendre polynomial evaluation

    x: input values between 0 and 1.
       Zero represents the present.
       One represents the timescale **in the past**, as far back as the system can remember.
    c: coefficient array
    """
    poly_len = c.shape[-1]
    batch_shape = c.shape[:-1]
    idx = np.arange(poly_len)

    eval_matrix = eval_legendre(idx[:, None], 2 * x - 1).T
    eval_matrix *= np.sqrt(2 * idx + 1) * (-1) ** idx
    eval_matrix = eval_matrix[None, ...]
    result = eval_matrix * c[..., None, :]
    return np.sum(result, axis=-1).reshape(batch_shape + (len(x),))


def fourier_val(x_vals: np.ndarray, coeffs: np.ndarray) -> np.ndarray:
    N = coeffs.shape[-1] // 2
    x_vals = x_vals[:, np.newaxis]
    m = np.arange(N)[None]
    two_pi_mx = 2 * np.pi * m * x_vals
    terms = np.sqrt(2.0) * np.exp(1j * two_pi_mx)
    terms[:, :1] /= np.sqrt(2.0)
    cos_terms = terms.real
    sin_terms = terms.imag

    cosine_coeffs = coeffs[..., 2 * m]
    sine_coeffs = coeffs[..., 2 * m + 1]
    fourier_sum = np.sum(cosine_coeffs * cos_terms + sine_coeffs * sin_terms, axis=-1)
    return fourier_sum


def legsval(x: np.ndarray, c: np.ndarray) -> np.ndarray:
    poly_len = c.shape[-1]
    batch_shape = c.shape[:-1]
    idx = np.arange(poly_len)

    eval_matrix = eval_legendre(idx[:, None], 1 - 2 * np.exp(-x)).T
    eval_matrix *= np.sqrt(2 * idx + 1) * (-1) ** idx
    eval_matrix = eval_matrix[None, ...]
    result = eval_matrix * c[..., None, :]
    return np.sum(result, axis=-1).reshape(batch_shape + (len(x),))


def _legendre_series_eval_single_z(z: jnp.ndarray, c: jnp.ndarray) -> jnp.ndarray:
    """
    Evaluate sum_k c[k] * sqrt(2k+1) * (-1)^k * P_k(z) for a single scalar z.

    z: scalar in [-1, 1] (ideally)
    c: shape (N,)
    returns: scalar
    """
    N = c.shape[0]

    k = jnp.arange(N, dtype=c.dtype)
    weights = jnp.sqrt(2.0 * k + 1.0) * jnp.power(-1.0, k)
    a = c * weights

    # N == 1
    def case_N1():
        return a[0]

    # N == 2
    def case_N2():
        return a[0] + a[1] * z

    # N >= 3
    def case_Nge3():
        P0 = jnp.ones_like(z)
        P1 = z
        out0 = a[0] * P0 + a[1] * P1

        # iterate k_int = 1..N-2 to build P_{k+1}
        ks = jnp.arange(1, N - 1, dtype=jnp.int32)

        def step(carry, k_int):
            Pkm1, Pk, out = carry
            kf = k_int.astype(c.dtype)  # float for recurrence math
            Pkp1 = ((2.0 * kf + 1.0) * z * Pk - kf * Pkm1) / (kf + 1.0)
            out = out + a[k_int + 1] * Pkp1
            return (Pk, Pkp1, out), None

        (_, _, out), _ = lax.scan(step, (P0, P1, out0), ks)
        return out

    return lax.cond(
        N == 1,
        case_N1,
        lambda: lax.cond(N == 2, case_N2, case_Nge3),
    )

# vmap over z, and over batch dims of c
_eval_over_z = jax.vmap(_legendre_series_eval_single_z, in_axes=(0, None), out_axes=0)

def legval_jax(x: jnp.ndarray, c: jnp.ndarray) -> jnp.ndarray:
    """
    JAX version of :func:`legval`.

    x: shape (T,), ideally in [0,1]
    c: shape (..., N)
    returns: shape (..., T)
    """
    # Map x in [0,1] to z in [-1,1]
    z = 2.0 * x - 1.0

    # Flatten batch dims so we can vmap cleanly
    batch_shape = c.shape[:-1]
    N = c.shape[-1]
    c2 = c.reshape((-1, N))  # (B, N)

    def eval_one_coeff_row(ci):
        return _eval_over_z(z, ci)  # (T,)

    y2 = jax.vmap(eval_one_coeff_row, in_axes=0, out_axes=0)(c2)  # (B, T)
    return y2.reshape(batch_shape + (x.shape[0],))

def legsval_jax(x: jnp.ndarray, c: jnp.ndarray) -> jnp.ndarray:
    """
    JAX version of :func:`legsval`.

    x: shape (T,), typically x>=0
    c: shape (..., N)
    returns: shape (..., T)
    """
    # Stable computation of z = 1 - 2*exp(-x)
    # exp(-x) = 1 + expm1(-x) => z = 1 - 2*(1 + expm1(-x)) = -1 - 2*expm1(-x)
    z = -1.0 - 2.0 * jnp.expm1(-x)

    batch_shape = c.shape[:-1]
    N = c.shape[-1]
    c2 = c.reshape((-1, N))  # (B, N)

    def eval_one_coeff_row(ci):
        return _eval_over_z(z, ci)  # (T,)

    y2 = jax.vmap(eval_one_coeff_row, in_axes=0, out_axes=0)(c2)  # (B, T)
    return y2.reshape(batch_shape + (x.shape[0],))



__all__ = [
    "EvalFn",
    "SystemParams",
    "fourier_val",
    "get_L_vector",
    "get_output_vector",
    "get_system_params",
    "get_implicit_M",
    "get_implicit_legt_A",
    "legval",
    "legsval",
    "legval_jax",
    "legsval_jax",
    "make_m_mat",
]


if __name__ == "__main__":
    # -----------------------
    # Config
    # -----------------------
    key = jax.random.PRNGKey(0)
    T = 50          # number of evaluation points
    N = 32          # polynomial order
    batch_shape = (4, 3)  # arbitrary batch dims

    # -----------------------
    # Generate test inputs
    # -----------------------
    # x in [0,1] for legval
    key, sub = jax.random.split(key)
    x_leg = jax.random.uniform(sub, (T,), minval=0.0, maxval=1.0)

    # x >= 0 for legsval
    key, sub = jax.random.split(key)
    x_legs = jax.random.exponential(sub, (T,))

    # coefficients
    key, sub = jax.random.split(key)
    c = jax.random.normal(sub, batch_shape + (N,))

    # Convert to NumPy for reference implementations
    x_leg_np = np.array(x_leg)
    x_legs_np = np.array(x_legs)
    c_np = np.array(c)

    # -----------------------
    # Evaluate
    # -----------------------
    y_leg_np = legval(x_leg_np, c_np)
    y_leg_jax = legval_jax(x_leg, c)

    y_legs_np = legsval(x_legs_np, c_np)
    y_legs_jax = legsval_jax(x_legs, c)

    # -----------------------
    # Errors
    # -----------------------
    leg_abs_err = np.max(np.abs(y_leg_np - np.array(y_leg_jax)))
    leg_rel_err = leg_abs_err / (np.max(np.abs(y_leg_np)) + 1e-12)

    legs_abs_err = np.max(np.abs(y_legs_np - np.array(y_legs_jax)))
    legs_rel_err = legs_abs_err / (np.max(np.abs(y_legs_np)) + 1e-12)

    # -----------------------
    # Report
    # -----------------------
    print("legval vs legval_jax")
    print(f"  max abs error: {leg_abs_err:.3e}")
    print(f"  max rel error: {leg_rel_err:.3e}")

    print("\nlegsval vs legsval_jax")
    print(f"  max abs error: {legs_abs_err:.3e}")
    print(f"  max rel error: {legs_rel_err:.3e}")