File size: 3,849 Bytes
b0e88cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | # EVOLVE-BLOCK-START
import jax
import jax.numpy as jnp
import optax
import numpy as np
from dataclasses import dataclass
@dataclass
class Hyperparameters:
"""Hyperparameters for the optimization process."""
num_intervals: int = 50
learning_rate: float = 0.01
num_steps: int = 15000
warmup_steps: int = 1000
class C2Optimizer:
"""
Optimizes a discretized function to find a lower bound for the C2 constant
using the rigorous, unitless, piecewise-linear integral method.
"""
def __init__(self, hypers: Hyperparameters):
self.hypers = hypers
def _objective_fn(self, f_values: jnp.ndarray) -> jnp.ndarray:
"""
Computes the objective function using the unitless norm calculation.
"""
f_non_negative = jax.nn.relu(f_values)
# Unscaled discrete autoconvolution
N = self.hypers.num_intervals
padded_f = jnp.pad(f_non_negative, (0, N))
fft_f = jnp.fft.fft(padded_f)
convolution = jnp.fft.ifft(fft_f * fft_f).real
# Calculate L2-norm squared of the convolution (rigorous method)
num_conv_points = len(convolution)
h = 1.0 / (num_conv_points + 1)
y_points = jnp.concatenate([jnp.array([0.0]), convolution, jnp.array([0.0])])
y1, y2 = y_points[:-1], y_points[1:]
l2_norm_squared = jnp.sum((h / 3) * (y1**2 + y1 * y2 + y2**2))
# Calculate L1-norm of the convolution
norm_1 = jnp.sum(jnp.abs(convolution)) / (len(convolution) + 1)
# Calculate infinity-norm of the convolution
norm_inf = jnp.max(jnp.abs(convolution))
# Calculate C2 ratio
denominator = norm_1 * norm_inf
c2_ratio = l2_norm_squared / denominator
# We want to MAXIMIZE C2, so the optimizer must MINIMIZE its negative.
return -c2_ratio
def train_step(self, f_values: jnp.ndarray, opt_state: optax.OptState) -> tuple:
"""Performs a single training step."""
loss, grads = jax.value_and_grad(self._objective_fn)(f_values)
updates, opt_state = self.optimizer.update(grads, opt_state, f_values)
f_values = optax.apply_updates(f_values, updates)
return f_values, opt_state, loss
def run_optimization(self):
"""Sets up and runs the full optimization process."""
schedule = optax.warmup_cosine_decay_schedule(
init_value=0.0,
peak_value=self.hypers.learning_rate,
warmup_steps=self.hypers.warmup_steps,
decay_steps=self.hypers.num_steps - self.hypers.warmup_steps,
end_value=self.hypers.learning_rate * 1e-4,
)
self.optimizer = optax.adam(learning_rate=schedule)
key = jax.random.PRNGKey(42)
f_values = jax.random.uniform(key, (self.hypers.num_intervals,))
opt_state = self.optimizer.init(f_values)
print(
f"Number of intervals (N): {self.hypers.num_intervals}, Steps: {self.hypers.num_steps}"
)
train_step_jit = jax.jit(self.train_step)
loss = jnp.inf
for step in range(self.hypers.num_steps):
f_values, opt_state, loss = train_step_jit(f_values, opt_state)
if step % 1000 == 0 or step == self.hypers.num_steps - 1:
print(f"Step {step:5d} | C2 ≈ {-loss:.8f}")
final_c2 = -self._objective_fn(f_values)
print(f"Final C2 lower bound found: {final_c2:.8f}")
return jax.nn.relu(f_values), final_c2
def run():
"""Entry point for running the optimization."""
hypers = Hyperparameters()
optimizer = C2Optimizer(hypers)
optimized_f, final_c2_val = optimizer.run_optimization()
loss_val = -final_c2_val
f_values_np = np.array(optimized_f)
return f_values_np, float(final_c2_val), float(loss_val), hypers.num_intervals
# EVOLVE-BLOCK-END
|