| """ |
| 1. Run seed 100 at N=4000, save params |
| 2. Upsample to N=8000, refine |
| 3. Perturbation restarts from best |
| 4. Also scan seeds 80-130 for better starts |
| """ |
| import sys |
| 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 optimize(N, seed=None, init_params_np=None, adam_steps=80000, lr=0.005, temp=300.0): |
| dx = 0.5 / N |
|
|
| @jax.jit |
| def obj_smooth(params, t): |
| f = jnp.exp(jnp.clip(params, -8, 4)) |
| 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 |
| integral_sq = (jnp.sum(f) * dx) ** 2 |
| smooth_max = jax.nn.logsumexp(t * conv) / t |
| return smooth_max / integral_sq |
|
|
| @jax.jit |
| def obj_hard(params): |
| f = jnp.exp(jnp.clip(params, -8, 4)) |
| 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 |
| integral_sq = (jnp.sum(f) * dx) ** 2 |
| return jnp.max(conv) / integral_sq |
|
|
| grad_smooth = jax.jit(jax.grad(obj_smooth)) |
| grad_hard = jax.jit(jax.grad(obj_hard)) |
|
|
| if init_params_np is not None: |
| if len(init_params_np) != N: |
| |
| old_f = np.exp(np.clip(init_params_np, -8, 4)) |
| new_f = np.interp(np.linspace(0, 1, N), np.linspace(0, 1, len(init_params_np)), old_f) |
| params = jnp.array(np.log(np.maximum(new_f, 1e-6))) |
| else: |
| params = jnp.array(init_params_np) |
| else: |
| np.random.seed(seed) |
| init_f = np.ones(N) * 0.5 + 0.02 * np.random.randn(N) |
| params = jnp.array(np.log(np.maximum(init_f, 1e-6))) |
|
|
| lr_schedule = optax.warmup_cosine_decay_schedule( |
| init_value=0.0, peak_value=lr, warmup_steps=min(2000, adam_steps//5), |
| decay_steps=adam_steps - min(2000, adam_steps//5), end_value=lr * 1e-4, |
| ) |
| optimizer = optax.adam(learning_rate=lr_schedule) |
| opt_state = optimizer.init(params) |
|
|
| best_c1 = float('inf') |
| best_params = params |
|
|
| for step in range(adam_steps): |
| loss, grads = jax.value_and_grad(obj_smooth)(params, temp) |
| updates, opt_state = optimizer.update(grads, opt_state, params) |
| params = optax.apply_updates(params, updates) |
|
|
| if step >= adam_steps - 2000: |
| hc = float(obj_hard(params)) |
| if hc < best_c1: |
| best_c1 = hc |
| best_params = params |
|
|
| |
| params_np = np.array(best_params, dtype=np.float64) |
| for t in [1000.0, 10000.0]: |
| def scipy_obj(p): |
| p_jax = jnp.array(p) |
| val = float(obj_smooth(p_jax, t)) |
| g = np.array(grad_smooth(p_jax, t), 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-14, 'maxcor': 100}, |
| ) |
| params_np = result.x |
|
|
| for _ in range(3): |
| def scipy_hard(p): |
| p_jax = jnp.array(p) |
| val = float(obj_hard(p_jax)) |
| g = np.array(grad_hard(p_jax), dtype=np.float64) |
| return val, g |
| result = scipy_minimize( |
| scipy_hard, params_np, method='L-BFGS-B', jac=True, |
| options={'maxiter': 20000, 'ftol': 1e-16, 'gtol': 1e-15, 'maxcor': 100}, |
| ) |
| params_np = result.x |
|
|
| f_final = np.exp(np.clip(params_np, -8, 4)) |
| c1_final = compute_c1_numpy(f_final, N) |
| return params_np, f_final, c1_final |
|
|
|
|
| def run(): |
| N = 4000 |
| best_c1_overall = float('inf') |
| best_f_overall = None |
| best_params_overall = None |
|
|
| |
| sys.stdout.write("Phase 1: Scanning seeds at N=4000\n") |
| sys.stdout.flush() |
|
|
| for seed in range(80, 131): |
| params, f, c1 = optimize(N, seed=seed, adam_steps=60000) |
| sys.stdout.write(f" Seed {seed:3d}: C1={c1:.10f}") |
| if c1 < best_c1_overall: |
| best_c1_overall = c1 |
| best_f_overall = f |
| best_params_overall = params |
| sys.stdout.write(" ***") |
| sys.stdout.write("\n") |
| sys.stdout.flush() |
|
|
| sys.stdout.write(f"\nBest after scan: C1={best_c1_overall:.10f}\n") |
|
|
| |
| sys.stdout.write("Phase 2: Upsample to N=8000\n") |
| sys.stdout.flush() |
| N2 = 8000 |
| params2, f2, c1_2 = optimize(N2, init_params_np=best_params_overall, |
| adam_steps=40000, lr=0.002) |
| sys.stdout.write(f" N=8000: C1={c1_2:.10f}\n") |
| sys.stdout.flush() |
|
|
| if c1_2 < best_c1_overall: |
| best_c1_overall = c1_2 |
| best_f_overall = f2 |
| best_params_overall = params2 |
| N = N2 |
|
|
| |
| sys.stdout.write("Phase 3: Perturbation restarts\n") |
| sys.stdout.flush() |
| for i in range(5): |
| key = jax.random.PRNGKey(i * 17) |
| noise = 0.03 * jax.random.normal(key, shape=(len(best_params_overall),)) |
| perturbed = best_params_overall + np.array(noise) |
| p, f, c1 = optimize(len(best_f_overall), init_params_np=perturbed, |
| adam_steps=30000, lr=0.002) |
| sys.stdout.write(f" Perturb {i}: C1={c1:.10f}") |
| if c1 < best_c1_overall: |
| best_c1_overall = c1 |
| best_f_overall = f |
| best_params_overall = p |
| N = len(f) |
| sys.stdout.write(" ***") |
| sys.stdout.write("\n") |
| sys.stdout.flush() |
|
|
| sys.stdout.write(f"\nFinal C1: {best_c1_overall:.10f}\n") |
| sys.stdout.flush() |
| return best_f_overall, best_c1_overall, best_c1_overall, len(best_f_overall) |
|
|
