def penalty(constraint, alpha=0.1, sigma=5, transition_width=0.2):
"""
Penalty function using tanh for smooth transition
Parameters:
- constraint: constraint value g
- alpha: penalty parameter
- sigma: switching point
- transition_width: transition width, smaller values create steeper transitions
"""
def safe_log(x):
x = jnp.clip(x, 1e-10, 1e6)
return jnp.log(x)
# Function values for the two branches
quadratic_barrier = alpha / 2 * (jnp.square((constraint - 2 * sigma) / sigma) - jnp.ones_like(constraint))
log_barrier = -alpha * safe_log(constraint)
combined_barrier = quadratic_barrier + log_barrier # Left-side combination
# Smooth transition weight (between 0 and 1)
# When g << σ, weight ≈ 0 (use combined_barrier)
# When g >> σ, weight ≈ 1 (use log_barrier)
weight = 0.5 * (1 + jnp.tanh((constraint - sigma) / transition_width))
# Smooth interpolation
smooth_result = weight * log_barrier + (1 - weight) * combined_barrier
return jnp.clip(smooth_result, 0, 1e8)
