ccevolve/tx_workspace/ac1_ours/offline_optimize.py

"""Long-running offline optimization to find the best sequence."""
import time
import numpy as np
from scipy.signal import fftconvolve
from scipy.optimize import minimize

def actual_obj(a):
    a = np.maximum(a, 0)
    n = len(a)
    conv = fftconvolve(a, a)
    s = np.sum(a)
    if s < 0.01: return float('inf')
    return 2.0 * n * np.max(conv) / s**2

def make_lp_obj_grad(n, p):
    def f(a):
        a = np.maximum(a, 1e-12)
        S = np.sum(a)
        conv = fftconvolve(a, a)
        conv = np.maximum(conv, 1e-30)
        log_conv = np.log(conv)
        lcm = np.max(log_conv)
        lcs = log_conv - lcm
        exp_p_lcs = np.exp(p * lcs)
        sum_exp = np.sum(exp_p_lcs)
        Lp = np.exp(lcm) * sum_exp ** (1.0/p)
        obj = 2.0 * n * Lp / S**2
        w = (sum_exp ** ((1-p)/p)) * np.exp((p-1) * lcs)
        G = fftconvolve(w, a[::-1], mode='valid')
        G = G[:n] if len(G) >= n else np.pad(G, (0, n-len(G)))
        dLp_da = 2 * G
        dobj_da = 2.0 * n / S**2 * (dLp_da - 2.0 * Lp / S)
        return obj, dobj_da
    return f

def make_lse_obj_grad(n, beta):
    def f(a):
        a = np.maximum(a, 1e-12)
        S = np.sum(a)
        conv = fftconvolve(a, a)
        conv_max = np.max(conv)
        shifted = beta * (conv - conv_max)
        exp_shifted = np.exp(shifted)
        sum_exp = np.sum(exp_shifted)
        lse = conv_max + np.log(sum_exp) / beta
        obj = 2.0 * n * lse / S**2
        softmax_w = exp_shifted / sum_exp
        G = fftconvolve(softmax_w, a[::-1], mode='valid')
        G = G[:n] if len(G) >= n else np.pad(G, (0, n-len(G)))
        dlse_da = 2 * G
        dobj_da = 2.0 * n / S**2 * (dlse_da - 2.0 * lse / S)
        return obj, dobj_da
    return f

def deep_optimize(a0, time_budget=30.0):
    n = len(a0)
    bounds = [(1e-10, 1000.0)] * n
    t0 = time.time()
    best_a = a0.copy()
    best_val = actual_obj(a0)
    
    # Phase 1: Lp progression
    for p in [16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]:
        if time.time() - t0 > time_budget * 0.7:
            break
        remaining = time_budget * 0.7 - (time.time() - t0)
        maxiter = max(100, int(remaining * 500))
        res = minimize(make_lp_obj_grad(n, p), a0, method='L-BFGS-B',
                      jac=True, bounds=bounds,
                      options={'maxiter': maxiter, 'ftol': 1e-15, 'gtol': 1e-14})
        a0 = np.maximum(res.x, 1e-10)
        val = actual_obj(a0)
        if val < best_val:
            best_val = val
            best_a = a0.copy()
    
    # Phase 2: LSE refinement  
    for beta in [100, 300, 1000, 3000, 10000]:
        if time.time() - t0 > time_budget * 0.9:
            break
        remaining = time_budget * 0.9 - (time.time() - t0)
        maxiter = max(100, int(remaining * 500))
        res = minimize(make_lse_obj_grad(n, beta), a0, method='L-BFGS-B',
                      jac=True, bounds=bounds,
                      options={'maxiter': maxiter, 'ftol': 1e-15, 'gtol': 1e-14})
        a0 = np.maximum(res.x, 1e-10)
        val = actual_obj(a0)
        if val < best_val:
            best_val = val
            best_a = a0.copy()
    
    # Phase 3: Final Lp polish
    for p in [2048, 8192]:
        if time.time() - t0 > time_budget - 0.5:
            break
        remaining = time_budget - (time.time() - t0)
        maxiter = max(100, int(remaining * 500))
        res = minimize(make_lp_obj_grad(n, p), best_a, method='L-BFGS-B',
                      jac=True, bounds=bounds,
                      options={'maxiter': maxiter, 'ftol': 1e-15, 'gtol': 1e-14})
        a_test = np.maximum(res.x, 1e-10)
        val = actual_obj(a_test)
        if val < best_val:
            best_val = val
            best_a = a_test.copy()
    
    return best_a, best_val

rng = np.random.default_rng(42)

best_val = float('inf')
best_a = None
best_n = None

# Try many starting points for n=150-300
t_total = time.time()
for n in [150, 200, 250, 300]:
    for trial in range(30):
        # Diverse starting points
        kind = rng.integers(0, 6)
        if kind == 0:
            a0 = rng.exponential(rng.uniform(0.5, 5), n)
        elif kind == 1:
            a0 = rng.uniform(0.1, 10, n)
        elif kind == 2:
            a0 = rng.gamma(rng.uniform(0.5, 3), rng.uniform(0.5, 3), n)
        elif kind == 3:
            a0 = np.abs(rng.standard_cauchy(n)) + 0.1
        elif kind == 4:
            a0 = rng.chisquare(rng.uniform(1, 5), n)
        else:
            a0 = rng.lognormal(rng.uniform(-1, 2), rng.uniform(0.3, 2), n)
        
        a_opt, val = deep_optimize(a0.copy(), time_budget=5.0)
        
        if val < best_val:
            best_val = val
            best_a = a_opt.copy()
            best_n = n
            print(f"n={n}, trial={trial}: NEW BEST {val:.6f}")

print(f"\nBest overall: {best_val:.6f} at n={best_n}")
print(f"Total time: {time.time()-t_total:.1f}s")

# Deep refinement of the best
print("\nDeep refinement of best...")
a_final, v_final = deep_optimize(best_a.copy(), time_budget=60.0)
print(f"After deep refinement: {v_final:.6f}")

# Save the best sequence
np.save('/workspace/best_sequence.npy', a_final)
print(f"Saved best sequence (n={len(a_final)})")