- register2 passing
This commit is contained in:
w.pomp
+40
-53
@@ -188,7 +188,9 @@ pub struct AsgdConfig {
|
||||
pub max_iterations: usize,
|
||||
/// Convergence tolerance (gradient norm)
|
||||
pub tolerance: f64,
|
||||
/// Maximum step length (displacement in mm). Default: 1.0
|
||||
/// Maximum step length. In elastix this is `delta` — used both for initial
|
||||
/// learning rate estimation (`a = delta * (A+1)^alpha / (jacg + eps)`) and
|
||||
/// as the per-iteration step clamp. Default: 1.0
|
||||
pub maximum_step_length: f64,
|
||||
/// Gain parameter A (denominator offset). Default: 20.0
|
||||
pub sp_a: f64,
|
||||
@@ -298,21 +300,25 @@ where
|
||||
|
||||
// Compute learning rate: a(t_k) = a / (A + t_k + 1)^alpha
|
||||
let t_k: T = current_time;
|
||||
let a_t = a / (T::from(config.sp_a).unwrap() + t_k + T::one()).powf(alpha);
|
||||
let max_step: T = num::cast(config.maximum_step_length).unwrap();
|
||||
let a_t = T::min(
|
||||
a / (T::from(config.sp_a).unwrap() + t_k + T::one()).powf(alpha),
|
||||
max_step,
|
||||
);
|
||||
|
||||
// Update time using sigmoid: t_{k+1} = max(0, t_k + sigmoid(-g_k^T * g_{k-1}))
|
||||
// In elastix, the dot product is computed in SCALED space: g_scaled = g/scales.
|
||||
// In elastix, the dot product is computed in SCALED space: g_scaled = g/sqrt(C).
|
||||
let dot_product_val: f64 = if let Some(ref prev_g) = previous_gradient {
|
||||
let dot_product: T = if let Some(ref scales) = config.scales {
|
||||
// Scaled dot product: sum(g_i * prev_g_i / scales_i²)
|
||||
gradient
|
||||
.iter()
|
||||
.zip(prev_g.iter())
|
||||
.zip(scales.iter())
|
||||
.fold(T::zero(), |acc, ((&g, &pg), &s)| {
|
||||
// Scaled dot product: sum((g/sqrt(C)) * (prev_g/sqrt(C))) = sum(g*pg/C)
|
||||
gradient.iter().zip(prev_g.iter()).zip(scales.iter()).fold(
|
||||
T::zero(),
|
||||
|acc, ((&g, &pg), &s)| {
|
||||
let s2: T = num::cast(s.max(1e-10)).unwrap();
|
||||
acc + g * pg / (s2 * s2)
|
||||
})
|
||||
let s4 = s2 * s2;
|
||||
acc + g * pg / s4
|
||||
},
|
||||
)
|
||||
} else {
|
||||
gradient
|
||||
.iter()
|
||||
@@ -329,17 +335,17 @@ where
|
||||
};
|
||||
|
||||
// Elastix gradient descent step in unscaled parameter space:
|
||||
// Optimizer stores p_scaled = p * scales, updates p_scaled -= a_t * g_scaled
|
||||
// where g_scaled = g / scales (chain rule from GetDerivative).
|
||||
// Convert back: p_new = p_scaled_new / scales = p - a_t * g / scales²
|
||||
// So direction = -g / scales².
|
||||
// m_Position stores the unscaled position p.
|
||||
// GetScaledDerivative computes g_scaled = g / scales.
|
||||
// Update: p -= a_t * g_scaled = p - a_t * g / scales.
|
||||
// So direction = -g / scales.
|
||||
let direction: Vec<T> = match &config.scales {
|
||||
Some(scales) => gradient
|
||||
.iter()
|
||||
.zip(scales.iter())
|
||||
.map(|(&g, &s)| {
|
||||
let s2: T = num::cast(s.max(1e-10)).unwrap();
|
||||
-g / (s2 * s2)
|
||||
-g / s2
|
||||
})
|
||||
.collect(),
|
||||
None => gradient.iter().map(|&g| -g).collect(),
|
||||
@@ -409,14 +415,12 @@ fn sigmoid<T: Float + Debug>(x: T, fmax: T, fmin: T, omega: T) -> T {
|
||||
(fmax - fmin) * sigmoid_raw + fmin
|
||||
}
|
||||
|
||||
/// Estimate ASGD parameters from gradient statistics.
|
||||
/// Estimate ASGD parameters using DisplacementDistribution method.
|
||||
///
|
||||
/// Matches elastix `AutomaticParameterEstimationOriginal()`:
|
||||
/// - `a_max = A * delta / sigma1 / sqrt(maxJCJ)`
|
||||
/// - `sigma1 = sqrt(gg / TrC)` where `gg = ||g_scaled||^2`
|
||||
/// - `TrC = sum(C_scaled[i][i]) = n` (when scales = sqrt(C[i][i]), C_scaled[i][i] = 1)
|
||||
/// - `maxJCJ` is approximated as max of squared scaled Jacobian column norms
|
||||
/// - `g_scaled = g / scales` (golden standard normalization)
|
||||
/// Matches elastix `AutomaticParameterEstimationUsingDisplacementDistribution()`:
|
||||
/// - `a = delta * (A+1)^alpha / (jacg + eps)` where `delta = MaximumStepLength`
|
||||
/// - `jacg ≈ sqrt(gg) * 1.8` as analytical approximation
|
||||
/// - `maximum_step_length` is also the per-iteration step clamp (same value as delta).
|
||||
fn estimate_asgd_parameters<T, F>(
|
||||
_f: &F,
|
||||
_initial_point: &[T],
|
||||
@@ -432,54 +436,37 @@ where
|
||||
let fmax = T::from(config.sigmoid_max).unwrap();
|
||||
let delta = T::from(config.maximum_step_length).unwrap();
|
||||
|
||||
// Compute scaled gradient: g_scaled = g / scales
|
||||
// And its squared norm: gg = ||g/scales||^2
|
||||
let (gg, trc, max_jcj) = if let Some(ref scales) = config.scales {
|
||||
// Compute gg = ||g/scales||^2 = sum(g²/scales²)
|
||||
let gg = if let Some(ref scales) = config.scales {
|
||||
let mut gg = T::zero();
|
||||
for (&g, &s) in initial_gradient.iter().zip(scales.iter()) {
|
||||
let s_t = T::from(s.max(1e-10)).unwrap();
|
||||
let g_scaled = g / s_t;
|
||||
gg = gg + g_scaled * g_scaled;
|
||||
}
|
||||
|
||||
// Compute maxJCJ analytically for golden standard scales.
|
||||
// Elastix: maxJCJ = max_j [Tr(J_j C J_j^T) + 2√2 ||J_j C J_j^T||_F]
|
||||
// For golden standard scales, C ≈ diag(scales²), so C_scaled ≈ I.
|
||||
// Then J_scaled J_scaled^T = diag(a, ..., a) where
|
||||
// a = Σ_j max_u((u-c)²)/s_j² + 1 = Σ_j 3(N_j-1)/(N_j+1) + 1
|
||||
// and maxJCJ = a · (ndim + 2√(2·ndim)).
|
||||
let n_params = scales.len();
|
||||
let ndim = ((-1.0 + (1.0 + 4.0 * n_params as f64).sqrt()) / 2.0).round() as usize;
|
||||
let mut a_val = 1.0; // translation contribution: 1/s² = 1
|
||||
for j in 0..ndim {
|
||||
let s_j = scales[j];
|
||||
let n_j = (12.0 * s_j * s_j + 1.0).sqrt();
|
||||
a_val += 3.0 * (n_j - 1.0) / (n_j + 1.0);
|
||||
}
|
||||
let max_jcj = a_val * (ndim as f64 + 2.0 * (2.0 * ndim as f64).sqrt());
|
||||
|
||||
let n = T::from(n_params).unwrap();
|
||||
(gg, n, T::from(max_jcj).unwrap())
|
||||
gg
|
||||
} else {
|
||||
let mut gg = T::zero();
|
||||
for &g in initial_gradient.iter() {
|
||||
gg = gg + g * g;
|
||||
}
|
||||
let n = T::from(initial_gradient.len()).unwrap();
|
||||
(gg, n, T::one())
|
||||
gg
|
||||
};
|
||||
|
||||
// Elastix: sigma1 = sqrt(gg / TrC), a_max = A * delta / sigma1 / sqrt(maxJCJ)
|
||||
let sigma1 = if gg > T::from(1e-14).unwrap() && trc > T::from(1e-14).unwrap() {
|
||||
(gg / trc).sqrt()
|
||||
// Compute jacg = displacement distribution estimate.
|
||||
// For golden standard scales: E[||d||^2] = gg, so jacg ≈ 1.8 * sqrt(gg)
|
||||
let sqrt_gg = if gg > T::from(1e-20).unwrap() {
|
||||
gg.sqrt()
|
||||
} else {
|
||||
T::zero()
|
||||
};
|
||||
let jacg = sqrt_gg * T::from(1.8).unwrap();
|
||||
|
||||
let a = if sigma1 > T::from(1e-14).unwrap() && max_jcj > T::from(1e-14).unwrap() {
|
||||
a_param * delta / sigma1 / max_jcj.sqrt()
|
||||
// Elastix DisplacementDistribution: a = delta * (A+1)^alpha / (jacg + eps)
|
||||
let a = if jacg > T::from(1e-14).unwrap() {
|
||||
delta * a_param.powf(alpha) / jacg
|
||||
} else {
|
||||
a_param
|
||||
delta * a_param.powf(alpha)
|
||||
};
|
||||
|
||||
let fmin = T::from(config.sigmoid_min).unwrap();
|
||||
|
||||
Reference in New Issue
Block a user