ccevolve/tx_workspace/ac1_alphaevolve/solution_v10.py

"""
Key innovations:
1. Normalize f so integral=1 (removes scale degeneracy)
2. Use Lp norm as smooth max approximation
3. Multi-phase: start with low p, increase gradually
4. Many diverse random restarts
"""
import jax
import jax.numpy as jnp
import numpy as np
from scipy.optimize import minimize as scipy_minimize
import optax


def compute_c1_numpy(f_values, n_points):
    dx = 0.5 / n_points
    f_nn = np.maximum(f_values, 0.0)
    autoconv = np.convolve(f_nn, f_nn, mode='full') * dx
    integral_sq = (np.sum(f_nn) * dx) ** 2
    if integral_sq < 1e-12:
        return 1e10
    return np.max(autoconv) / integral_sq


def make_fns(N):
    dx = 0.5 / N

    @jax.jit
    def params_to_f(params):
        """Convert params to normalized non-negative function"""
        f = jax.nn.softplus(params)
        # Normalize so integral = 1
        integral = jnp.sum(f) * dx
        f = f / jnp.maximum(integral, 1e-9)
        return f

    @jax.jit
    def compute_c1_smooth(params, p):
        """Lp norm approximation to C1"""
        f = params_to_f(params)
        padded = jnp.zeros(2 * N)
        padded = padded.at[:N].set(f)
        fft_f = jnp.fft.rfft(padded)
        conv = jnp.fft.irfft(fft_f * fft_f, n=2 * N) * dx
        # C1 = max(conv) since integral = 1
        # Approximate with Lp norm
        # (sum(conv^p) / (2N))^(1/p) -> max as p->inf
        conv_pos = jnp.maximum(conv, 0.0)
        lp = (jnp.mean(conv_pos ** p)) ** (1.0 / p)
        return lp

    @jax.jit
    def compute_c1_logsumexp(params, temp):
        """LogSumExp approximation to C1"""
        f = params_to_f(params)
        padded = jnp.zeros(2 * N)
        padded = padded.at[:N].set(f)
        fft_f = jnp.fft.rfft(padded)
        conv = jnp.fft.irfft(fft_f * fft_f, n=2 * N) * dx
        return jax.nn.logsumexp(temp * conv) / temp

    @jax.jit
    def compute_c1_hard(params):
        f = params_to_f(params)
        padded = jnp.zeros(2 * N)
        padded = padded.at[:N].set(f)
        fft_f = jnp.fft.rfft(padded)
        conv = jnp.fft.irfft(fft_f * fft_f, n=2 * N) * dx
        return jnp.max(conv)

    grad_lse = jax.jit(jax.grad(compute_c1_logsumexp))

    return params_to_f, compute_c1_smooth, compute_c1_logsumexp, compute_c1_hard, grad_lse


def optimize_single(N, seed, steps=80000, verbose=True):
    dx = 0.5 / N
    params_to_f, c1_smooth, c1_lse, c1_hard, grad_lse = make_fns(N)

    np.random.seed(seed)
    x = np.linspace(0, 1, N)

    # Diverse initializations
    init_types = [
        lambda: np.ones(N),  # constant
        lambda: np.exp(-10 * (x - 0.5)**2) + 0.1,  # centered Gaussian
        lambda: np.exp(-5 * (x - 0.3)**2) + 0.05,  # left Gaussian
        lambda: np.exp(-5 * (x - 0.7)**2) + 0.05,  # right Gaussian
        lambda: 0.5 + 0.5 * np.cos(2*np.pi*x),  # cosine
        lambda: np.maximum(1 - 4*np.abs(x - 0.5), 0) + 0.1,  # triangle
        lambda: np.where((x > 0.1) & (x < 0.9), 1.0, 0.1),  # wide box
        lambda: np.where((x > 0.2) & (x < 0.8), 1.0, 0.1),  # medium box
        lambda: np.exp(-50 * (x - 0.5)**2) + 0.01,  # narrow Gaussian
        lambda: np.abs(np.random.randn(N)) * 0.3 + 0.2,  # random
        lambda: 1 - 0.5 * np.abs(np.sin(3*np.pi*x)),  # wavy
        lambda: np.exp(-3*(x - 0.4)**2) + 0.5*np.exp(-3*(x-0.6)**2) + 0.05,  # bimodal
    ]

    init_f = init_types[seed % len(init_types)]()
    init_f = np.maximum(init_f, 0.01)
    # Inverse softplus
    params = jnp.array(np.log(np.expm1(np.maximum(init_f, 1e-3))))

    # Adam optimization with Lp norm, increasing p
    lr_schedule = optax.warmup_cosine_decay_schedule(
        init_value=0.0, peak_value=0.005, warmup_steps=2000,
        decay_steps=steps - 2000, end_value=1e-6,
    )
    optimizer = optax.adam(learning_rate=lr_schedule)
    opt_state = optimizer.init(params)

    best_c1 = float('inf')
    best_params = params

    for step in range(steps):
        progress = min(step / steps, 1.0)
        # Start with p=4, increase to p=40
        p = 4.0 + progress * 36.0

        loss, grads = jax.value_and_grad(c1_smooth)(params, p)
        updates, opt_state = optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)

        if step % 20000 == 0 or step == steps - 1:
            hc = float(c1_hard(params))
            if verbose:
                print(f"  [{seed:2d}] Step {step:6d} | C1={hc:.8f} | p={p:.1f}")
            if hc < best_c1:
                best_c1 = hc
                best_params = params

    # L-BFGS polish with logsumexp
    params_np = np.array(best_params, dtype=np.float64)
    for temp in [500.0, 2000.0, 10000.0]:
        def scipy_obj(p_arr):
            p_jax = jnp.array(p_arr)
            val = float(c1_lse(p_jax, temp))
            g = np.array(grad_lse(p_jax, temp), dtype=np.float64)
            return val, g

        result = scipy_minimize(
            scipy_obj, params_np, method='L-BFGS-B', jac=True,
            options={'maxiter': 3000, 'ftol': 1e-15, 'gtol': 1e-12},
        )
        params_np = result.x

    # Get final function values (unnormalized for evaluator)
    f_norm = np.log1p(np.exp(params_np))  # softplus
    f_norm = np.maximum(f_norm, 0.0)
    c1_final = compute_c1_numpy(f_norm, N)

    if verbose:
        print(f"  [{seed:2d}] Final: C1={c1_final:.10f}")

    return f_norm, c1_final


def run():
    N = 3000
    best_c1 = float('inf')
    best_f = None

    for seed in range(12):
        f, c1 = optimize_single(N, seed, steps=60000)
        if c1 < best_c1:
            best_c1 = c1
            best_f = f
            print(f"  *** GLOBAL BEST: C1={c1:.10f} (seed={seed})")

    # Take best and do extended polishing at higher N
    print(f"\nExtended polishing at N=5000...")
    N2 = 5000
    # Upsample
    f_up = np.interp(np.linspace(0, 1, N2), np.linspace(0, 1, len(best_f)), best_f)
    f_up2, c1_up = optimize_single(N2, 99, steps=40000)

    # Also polish the upsampled version
    dx2 = 0.5 / N2
    _, _, c1_lse2, c1_hard2, grad_lse2 = make_fns(N2)
    params_up = jnp.array(np.log(np.expm1(np.maximum(f_up, 1e-3))))

    params_np = np.array(params_up, dtype=np.float64)
    for temp in [500.0, 2000.0, 10000.0, 50000.0]:
        def scipy_obj(p_arr):
            p_jax = jnp.array(p_arr)
            val = float(c1_lse2(p_jax, temp))
            g = np.array(grad_lse2(p_jax, temp), dtype=np.float64)
            return val, g

        result = scipy_minimize(
            scipy_obj, params_np, method='L-BFGS-B', jac=True,
            options={'maxiter': 5000, 'ftol': 1e-15, 'gtol': 1e-12},
        )
        params_np = result.x

    f_final = np.log1p(np.exp(params_np))
    c1_final = compute_c1_numpy(f_final, N2)
    print(f"Upsampled polished: C1={c1_final:.10f}")

    if c1_final < best_c1:
        best_c1 = c1_final
        best_f = f_final
        N = N2

    print(f"\nFinal best C1: {best_c1:.10f}")
    return best_f, best_c1, best_c1, len(best_f)