File size: 3,751 Bytes
2facf1f | 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 112 113 114 115 116 117 | # 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 = 600
learning_rate: float = 0.005
end_lr_factor: float = 1e-4
num_steps: int = 40000
warmup_steps: int = 2000
class AutocorrelationOptimizer:
"""
Optimizes a discretized function to find the minimal C1 constant.
"""
def __init__(self, hypers: Hyperparameters):
self.hypers = hypers
self.domain_width = 0.5
self.dx = self.domain_width / self.hypers.num_intervals
def _objective_fn(self, f_values: jnp.ndarray) -> jnp.ndarray:
"""
Computes the objective function, which is the C1 ratio.
We minimize this ratio to find a tight upper bound.
"""
f_non_negative = jax.nn.relu(f_values)
integral_f = jnp.sum(f_non_negative) * self.dx
eps = 1e-9
integral_f_safe = jnp.maximum(integral_f, eps)
N = self.hypers.num_intervals
padded_f = jnp.pad(f_non_negative, (0, N))
fft_f = jnp.fft.fft(padded_f)
fft_conv = fft_f * fft_f
conv_f_f = jnp.fft.ifft(fft_conv).real
# Scale by dx.
scaled_conv_f_f = conv_f_f * self.dx
max_conv = jnp.max(scaled_conv_f_f)
c1_ratio = max_conv / (integral_f_safe**2)
# Return the value to be MINIMIZED.
return c1_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 * self.hypers.end_lr_factor,
)
self.optimizer = optax.adam(learning_rate=schedule)
key = jax.random.PRNGKey(42)
N = self.hypers.num_intervals
f_values = jnp.zeros((N,))
start_idx, end_idx = N // 4, 3 * N // 4
f_values = f_values.at[start_idx:end_idx].set(1.0)
f_values += 0.05 * jax.random.uniform(key, (N,))
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 # Initialize loss
for step in range(self.hypers.num_steps):
f_values, opt_state, loss = train_step_jit(f_values, opt_state)
if step % 2000 == 0 or step == self.hypers.num_steps - 1:
# CORRECTED PRINTING: Show the positive loss value directly.
print(f"Step {step:5d} | C1 ≈ {loss:.8f}")
print(f"Final C1 found: {loss:.8f}")
return jax.nn.relu(f_values), loss
def run():
"""Entry point for running the optimization and returning results."""
hypers = Hyperparameters()
optimizer = AutocorrelationOptimizer(hypers)
optimized_f, final_loss_val = optimizer.run_optimization()
final_c1 = float(final_loss_val)
f_values_np = np.array(optimized_f)
return f_values_np, final_c1, final_loss_val, hypers.num_intervals
# EVOLVE-BLOCK-END |