1use ndarray::Array1;
2use std::collections::HashSet;
3
4pub use gam_problem::{SeedConfig, SeedRiskProfile};
5use crate::estimate::EstimationError;
6use gam_problem::OrderedRhoBounds;
7
8fn add_seed_dedup(seeds: &mut Vec<Array1<f64>>, seen: &mut HashSet<Vec<u64>>, seed: Array1<f64>) {
9 let key: Vec<u64> = seed.iter().map(|&v| v.to_bits()).collect();
10 if seen.insert(key) {
11 seeds.push(seed);
12 }
13}
14
15fn safe_ln_pos(x: f64) -> Option<f64> {
16 if x.is_finite() && x > 0.0 {
17 Some(x.ln())
18 } else {
19 None
20 }
21}
22
23fn spde_rho_triplet_from_log_tau_log_kappa_nu(
24 log_tau: f64,
25 log_kappa: f64,
26 nu: f64,
27 bounds: OrderedRhoBounds,
28) -> Option<Array1<f64>> {
29 if !(nu.is_finite() && nu > 1.0) {
30 return None;
31 }
32 let logc0 = 0.0;
33 let logc1 = safe_ln_pos(nu)?;
34 let logc2 = safe_ln_pos(0.5 * nu * (nu - 1.0))?;
35 let rho0 = bounds.clamp(log_tau + logc0 + 2.0 * nu * log_kappa);
36 let rho1 = bounds.clamp(log_tau + logc1 + 2.0 * (nu - 1.0) * log_kappa);
37 let rho2 = bounds.clamp(log_tau + logc2 + 2.0 * (nu - 2.0) * log_kappa);
38 Some(Array1::from_vec(vec![rho0, rho1, rho2]))
39}
40
41fn add_spde_manifold_seeds(
42 seeds: &mut Vec<Array1<f64>>,
43 seen: &mut HashSet<Vec<u64>>,
44 bounds: OrderedRhoBounds,
45 heuristic_rhos: Option<&[f64]>,
46 primary: &Array1<f64>,
47) {
48 if primary.len() != 3 {
49 return;
50 }
51 let tau_anchors = [primary[2], 0.0, -2.0, 2.0];
53 let log_kappa_grid = [-2.0, -1.0, 0.0, 1.0, 2.0];
54 let nu_grid = [1.25, 1.5, 2.0, 2.5, 3.0, 4.0];
55 for &tau in &tau_anchors {
56 for &lk in &log_kappa_grid {
57 for &nu in &nu_grid {
58 if let Some(seed) = spde_rho_triplet_from_log_tau_log_kappa_nu(tau, lk, nu, bounds)
59 {
60 add_seed_dedup(seeds, seen, seed);
61 }
62 }
63 }
64 }
65
66 if let Some(vals) = heuristic_rhos
69 && vals.len() == 3
70 {
71 let l0 = vals[0].exp();
72 let l1 = vals[1].exp();
73 let l2 = vals[2].exp();
74 if l0.is_finite() && l1.is_finite() && l2.is_finite() && l0 > 1e-12 && l2 > 1e-12 {
75 let r = (l1 * l1) / (l0 * l2);
76 if r > 2.0 {
77 let nu = r / (r - 2.0);
78 let kappa2 = l1 / ((r - 2.0) * l2);
79 if nu.is_finite() && nu > 1.0 && kappa2.is_finite() && kappa2 > 0.0 {
80 let log_kappa = 0.5 * kappa2.ln();
81 let c2 = 0.5 * nu * (nu - 1.0);
82 if c2.is_finite() && c2 > 0.0 {
83 let log_tau = (l2 / (c2 * kappa2.powf(nu - 2.0))).max(1e-12).ln();
84 let local_nu = [nu, (nu - 0.3).max(1.05), nu + 0.3];
85 let local_tau = [log_tau, log_tau - 1.0, log_tau + 1.0];
86 let local_kappa = [log_kappa, log_kappa - 0.5, log_kappa + 0.5];
87 for &t in &local_tau {
88 for &lk in &local_kappa {
89 for &n in &local_nu {
90 if let Some(seed) =
91 spde_rho_triplet_from_log_tau_log_kappa_nu(t, lk, n, bounds)
92 {
93 add_seed_dedup(seeds, seen, seed);
94 }
95 }
96 }
97 }
98 }
99 }
100 }
101 }
102 }
103}
104
105fn add_first_order_fallback_seeds(
106 seeds: &mut Vec<Array1<f64>>,
107 seen: &mut HashSet<Vec<u64>>,
108 bounds: OrderedRhoBounds,
109 heuristic_rhos: Option<&[f64]>,
110) {
111 let rho2_floor = bounds.lower();
114 let default_log_kappa = [-2.0, -1.0, 0.0, 1.0];
115 let default_log_tau = [0.0, -2.0, 2.0];
116 for &t in &default_log_tau {
117 for &lk in &default_log_kappa {
118 let rho0 = bounds.clamp(t + 2.0 * lk);
119 let rho1 = bounds.clamp(t);
120 add_seed_dedup(seeds, seen, Array1::from_vec(vec![rho0, rho1, rho2_floor]));
121 }
122 }
123 if let Some(vals) = heuristic_rhos
124 && vals.len() == 3
125 && vals[0].is_finite()
126 && vals[1].is_finite()
127 {
128 let l0 = vals[0].exp();
129 let l1 = vals[1].exp();
130 let kappa2 = l0 / l1;
131 if kappa2.is_finite() && kappa2 > 0.0 {
132 let lk = 0.5 * kappa2.ln();
133 let t = vals[1];
134 let rho0 = bounds.clamp(t + 2.0 * lk);
135 let rho1 = bounds.clamp(t);
136 add_seed_dedup(seeds, seen, Array1::from_vec(vec![rho0, rho1, rho2_floor]));
137 }
138 }
139}
140
141fn add_nu2_reverse_manifold_seeds(
142 seeds: &mut Vec<Array1<f64>>,
143 seen: &mut HashSet<Vec<u64>>,
144 bounds: OrderedRhoBounds,
145 primary: &Array1<f64>,
146) {
147 if primary.len() != 3 {
148 return;
149 }
150 let ln_two = 2.0_f64.ln();
151 let tau_anchors = [primary[2], 0.0, -2.0, 2.0];
152 let log_kappa_grid = [-2.0, -1.0, 0.0, 1.0, 2.0];
153 for &tau_rho in &tau_anchors {
154 for &log_kappa in &log_kappa_grid {
155 let rho2 = bounds.clamp(tau_rho);
158 let rho1 = bounds.clamp(tau_rho + ln_two + 2.0 * log_kappa);
159 let rho0 = bounds.clamp(tau_rho + 4.0 * log_kappa);
160 add_seed_dedup(seeds, seen, Array1::from_vec(vec![rho0, rho1, rho2]));
161 }
162 }
163}
164
165fn halton(mut index: usize, base: usize) -> f64 {
166 let mut f = 1.0_f64;
167 let mut r = 0.0_f64;
168 while index > 0 {
169 f /= base as f64;
170 r += f * (index % base) as f64;
171 index /= base;
172 }
173 r
174}
175
176fn first_primes(n: usize) -> Vec<usize> {
177 let mut primes = Vec::with_capacity(n);
178 let mut x = 2usize;
179 while primes.len() < n {
180 let mut is_prime = true;
181 let mut d = 2usize;
182 while d * d <= x {
183 if x.is_multiple_of(d) {
184 is_prime = false;
185 break;
186 }
187 d += 1;
188 }
189 if is_prime {
190 primes.push(x);
191 }
192 x += 1;
193 }
194 primes
195}
196
197pub fn generate_rho_candidates(
198 num_penalties: usize,
199 heuristic_rhos: Option<&[f64]>,
200 config: &SeedConfig,
201) -> Result<Vec<Array1<f64>>, EstimationError> {
202 let mut seeds = Vec::new();
203 let mut seen: HashSet<Vec<u64>> = HashSet::new();
204
205 let bounds = OrderedRhoBounds::new(config.bounds.0, config.bounds.1)?;
213 let max_seeds = config.max_seeds.max(1);
214 let risk_shift = config.risk_profile.anchor_rho_shift();
215
216 if num_penalties == 0 {
217 add_seed_dedup(&mut seeds, &mut seen, Array1::<f64>::zeros(0));
218 return Ok(seeds);
219 }
220
221 let num_aux = config.num_auxiliary_trailing.min(num_penalties);
224 let num_smoothing = num_penalties - num_aux;
225 let aux_initial: Vec<f64> = if num_aux > 0 {
226 heuristic_rhos
227 .filter(|h| h.len() == num_penalties)
228 .map(|h| {
229 h[num_smoothing..]
230 .iter()
231 .copied()
232 .map(|v| bounds.clamp(v))
233 .collect()
234 })
235 .unwrap_or_else(|| vec![0.0; num_aux])
236 } else {
237 Vec::new()
238 };
239 let heuristic_rhovec: Option<Array1<f64>> = heuristic_rhos.and_then(|vals| {
240 if vals.len() == num_penalties {
241 Some(Array1::from_iter(
242 vals[..num_smoothing]
243 .iter()
244 .copied()
245 .map(|v| bounds.clamp(v))
246 .chain(
247 vals[num_smoothing..]
248 .iter()
249 .copied()
250 .map(|v| bounds.clamp(v)),
251 ),
252 ))
253 } else {
254 None
255 }
256 });
257
258 let primary = heuristic_rhovec.clone().unwrap_or_else(|| {
259 Array1::<f64>::from_elem(num_penalties, bounds.clamp(risk_shift))
260 });
261 add_seed_dedup(&mut seeds, &mut seen, primary.clone());
262 add_seed_dedup(&mut seeds, &mut seen, Array1::zeros(num_penalties));
264 match config.risk_profile {
268 SeedRiskProfile::Gaussian | SeedRiskProfile::GaussianLocationScale => {}
269 SeedRiskProfile::GeneralizedLinear | SeedRiskProfile::Survival => {
270 add_seed_dedup(
271 &mut seeds,
272 &mut seen,
273 Array1::from_elem(num_penalties, bounds.upper()),
274 );
275 }
276 }
277 if num_smoothing == 3 {
283 let smoothing_primary =
284 Array1::from_vec(primary.iter().take(num_smoothing).copied().collect());
285 let smoothing_heuristic_lambdas = heuristic_rhos.and_then(|vals| {
286 if vals.len() >= num_smoothing {
287 Some(&vals[..num_smoothing])
288 } else {
289 None
290 }
291 });
292 let mut spde_prefix_seeds = Vec::new();
293 let mut spde_prefix_seen: HashSet<Vec<u64>> = HashSet::new();
294 add_seed_dedup(
296 &mut spde_prefix_seeds,
297 &mut spde_prefix_seen,
298 Array1::from_vec(vec![primary[0], primary[1], bounds.lower()]),
299 );
300 add_nu2_reverse_manifold_seeds(
302 &mut spde_prefix_seeds,
303 &mut spde_prefix_seen,
304 bounds,
305 &smoothing_primary,
306 );
307 add_first_order_fallback_seeds(
308 &mut spde_prefix_seeds,
309 &mut spde_prefix_seen,
310 bounds,
311 smoothing_heuristic_lambdas,
312 );
313 add_spde_manifold_seeds(
314 &mut spde_prefix_seeds,
315 &mut spde_prefix_seen,
316 bounds,
317 smoothing_heuristic_lambdas,
318 &smoothing_primary,
319 );
320 for prefix_seed in spde_prefix_seeds {
321 let mut seed = Array1::<f64>::zeros(num_penalties);
322 for i in 0..num_smoothing {
323 seed[i] = prefix_seed[i];
324 }
325 for (i, &v) in aux_initial.iter().enumerate() {
326 seed[num_smoothing + i] = v;
327 }
328 add_seed_dedup(&mut seeds, &mut seen, seed);
329 }
330 }
331
332 for ¢er in config.risk_profile.baseline_centers() {
334 add_seed_dedup(
335 &mut seeds,
336 &mut seen,
337 Array1::from_elem(num_penalties, bounds.clamp(center)),
338 );
339 }
340
341 let dims_to_touch = num_penalties.min(12);
342 let step_base = if num_penalties <= 4 {
343 2.0
344 } else if num_penalties <= 12 {
345 2.5
346 } else {
347 3.0
348 };
349 let high_dim_cluster_threshold = 10usize;
350
351 if num_penalties >= high_dim_cluster_threshold {
352 let mut sorted_idx: Vec<usize> = (0..num_penalties).collect();
355 sorted_idx.sort_by(|&i, &j| primary[i].total_cmp(&primary[j]));
356
357 let cluster_size = (num_penalties / 3).max(1);
358 let small_end = cluster_size.min(num_penalties);
359 let large_start = num_penalties.saturating_sub(cluster_size);
360 let small_cluster = &sorted_idx[..small_end];
361 let large_cluster = &sorted_idx[large_start..];
362
363 let small_scale = step_base;
364 let large_scale = step_base + 0.75;
365
366 let mut conflict_a = primary.clone();
367 for &i in large_cluster {
368 conflict_a[i] = bounds.clamp(primary[i] + large_scale);
369 }
370 for &i in small_cluster {
371 conflict_a[i] = bounds.clamp(primary[i] - small_scale);
372 }
373 add_seed_dedup(&mut seeds, &mut seen, conflict_a);
374
375 let mut conflict_b = primary.clone();
376 for &i in large_cluster {
377 conflict_b[i] = bounds.clamp(primary[i] - large_scale);
378 }
379 for &i in small_cluster {
380 conflict_b[i] = bounds.clamp(primary[i] + small_scale);
381 }
382 add_seed_dedup(&mut seeds, &mut seen, conflict_b);
383
384 let mut heavy_up = primary.clone();
385 for &i in large_cluster {
386 heavy_up[i] = bounds.clamp(primary[i] + large_scale);
387 }
388 add_seed_dedup(&mut seeds, &mut seen, heavy_up);
389
390 let mut light_down = primary.clone();
391 for &i in small_cluster {
392 light_down[i] = bounds.clamp(primary[i] - small_scale);
393 }
394 add_seed_dedup(&mut seeds, &mut seen, light_down);
395 } else {
396 for i in 0..dims_to_touch {
398 let scale = step_base + 0.25 * primary[i].abs().min(8.0);
399 for dir in [-1.0, 1.0] {
400 let mut s = primary.clone();
401 s[i] = bounds.clamp(primary[i] + dir * scale);
402 add_seed_dedup(&mut seeds, &mut seen, s);
403 }
404 }
405
406 let pair_dims = num_penalties.min(6);
407 for i in 0..pair_dims {
408 for j in (i + 1)..pair_dims {
409 let mut s1 = primary.clone();
410 s1[i] = bounds.clamp(primary[i] + step_base);
411 s1[j] = bounds.clamp(primary[j] - step_base);
412 add_seed_dedup(&mut seeds, &mut seen, s1);
413
414 let mut s2 = primary.clone();
415 s2[i] = bounds.clamp(primary[i] - step_base);
416 s2[j] = bounds.clamp(primary[j] + step_base);
417 add_seed_dedup(&mut seeds, &mut seen, s2);
418 }
419 }
420 }
421
422 for &shift in config.risk_profile.global_shifts() {
439 let swept = primary.mapv(|v| bounds.clamp(v + shift));
440 add_seed_dedup(&mut seeds, &mut seen, swept);
441 }
442
443 if let Some(probe_rho) = config.over_smoothing_probe_rho {
451 let mut probe = primary.clone();
452 for j in 0..num_smoothing {
453 probe[j] = bounds.clamp(probe_rho);
454 }
455 add_seed_dedup(&mut seeds, &mut seen, probe);
456 }
457
458 let exploratory = max_seeds.saturating_sub(seeds.len()).min(8);
461 if exploratory > 0 {
462 let primes = first_primes(num_penalties.max(1));
463 let amp = config.risk_profile.exploratory_amplitude();
464 for t in 0..exploratory {
465 let mut s = primary.clone();
466 for i in 0..num_penalties {
467 let u = halton(t + 1, primes[i]); let centered = 2.0 * u - 1.0; s[i] = bounds.clamp(primary[i] + amp * centered);
470 }
471 add_seed_dedup(&mut seeds, &mut seen, s);
472 }
473 }
474
475 if num_aux > 0 {
481 for seed in &mut seeds {
482 for (i, &v) in aux_initial.iter().enumerate() {
483 seed[num_smoothing + i] = v;
484 }
485 }
486 let mut deduped = Vec::new();
487 let mut seen2: HashSet<Vec<u64>> = HashSet::new();
488 for seed in seeds {
489 let key: Vec<u64> = seed.iter().map(|&v| v.to_bits()).collect();
490 if seen2.insert(key) {
491 deduped.push(seed);
492 }
493 }
494 seeds = deduped;
495 }
496
497 if seeds.len() > max_seeds {
498 seeds.truncate(max_seeds);
499 }
500
501 if seeds.is_empty() {
502 seeds.push(Array1::<f64>::zeros(num_penalties));
503 }
504
505 Ok(seeds)
506}
507
508#[cfg(test)]
509mod tests {
510 use super::*;
511
512 #[test]
518 fn generate_rho_candidates_refuses_inverted_bounds() {
519 let cfg = SeedConfig {
520 bounds: (10.0, -10.0),
521 ..SeedConfig::default()
522 };
523 let err = generate_rho_candidates(3, None, &cfg)
524 .expect_err("an inverted seed box must be refused, not swapped");
525 assert!(
526 matches!(err, EstimationError::InvalidInput(_)),
527 "inverted seed box must be a typed InvalidInput, got {err:?}"
528 );
529 }
530
531 #[test]
535 fn generate_rho_candidates_accepts_ordered_bounds() {
536 let cfg = SeedConfig::default();
537 let seeds = generate_rho_candidates(3, None, &cfg).expect("ordered box is accepted");
538 assert!(!seeds.is_empty());
539 }
540
541 #[test]
542 fn uses_full_heuristicvector_as_primary_anchor() {
543 let cfg = SeedConfig {
544 risk_profile: SeedRiskProfile::Gaussian,
545 ..SeedConfig::default()
546 };
547 let heur = [-2.0, 0.0, 2.0];
548 let seeds = generate_rho_candidates(3, Some(&heur), &cfg).expect("ordered seed bounds");
549 assert!(!seeds.is_empty());
550 let first = &seeds[0];
551 assert_eq!(first.len(), 3);
552 assert!((first[0] - heur[0]).abs() < 1e-12);
553 assert!((first[1] - heur[1]).abs() < 1e-12);
554 assert!((first[2] - heur[2]).abs() < 1e-12);
555 }
556
557 #[test]
558 fn high_dim_uses_cluster_conflict_probeswithout_exploding() {
559 let cfg = SeedConfig {
560 max_seeds: 18,
561 risk_profile: SeedRiskProfile::GeneralizedLinear,
562 ..SeedConfig::default()
563 };
564 let heur = [-6.0, -5.0, -4.0, 0.0, 2.0, 4.0, -3.0, 0.0, 3.0, 5.0];
565 let seeds = generate_rho_candidates(10, Some(&heur), &cfg).expect("ordered seed bounds");
566 assert!(seeds.len() <= 18);
567 let primary = &seeds[0];
570 let has_conflict = seeds.iter().skip(1).any(|s| {
571 let mut any_up = false;
572 let mut any_down = false;
573 for i in 0..s.len() {
574 if s[i] > primary[i] {
575 any_up = true;
576 } else if s[i] < primary[i] {
577 any_down = true;
578 }
579 }
580 any_up && any_down
581 });
582 assert!(has_conflict);
583 }
584
585 #[test]
586 fn includes_neutralzero_seed() {
587 let cfg = SeedConfig::default();
588 let seeds = generate_rho_candidates(5, None, &cfg).expect("ordered seed bounds");
589 let haszero = seeds
590 .iter()
591 .any(|s| s.iter().all(|v| (*v - 0.0).abs() < 1e-12));
592 assert!(haszero);
593 }
594
595 #[test]
596 fn generalized_linear_seeds_include_early_stability_retreat_seed() {
597 let cfg = SeedConfig {
598 risk_profile: SeedRiskProfile::GeneralizedLinear,
599 ..SeedConfig::default()
600 };
601 let seeds = generate_rho_candidates(3, None, &cfg).expect("ordered seed bounds");
602 let retreat = Array1::from_elem(3, cfg.bounds.1);
603 let retreat_idx = seeds
604 .iter()
605 .position(|seed| seed == retreat)
606 .expect("generalized-linear seeds should include an upper-bound retreat seed");
607 assert!(
608 retreat_idx <= 2,
609 "retreat seed should be available before broader exploratory seeds: {retreat_idx}"
610 );
611 }
612
613 #[test]
614 fn three_penalty_seeds_include_nu2_reverse_manifold_triplets() {
615 let cfg = SeedConfig::default();
616 let seeds = generate_rho_candidates(3, None, &cfg).expect("ordered seed bounds");
617 let ln4 = 4.0_f64.ln();
618 let has_nu2_manifold_seed = seeds
619 .iter()
620 .any(|s| s.len() == 3 && ((2.0 * s[1] - s[0] - s[2]) - ln4).abs() < 1e-8);
621 assert!(has_nu2_manifold_seed);
622 }
623
624 #[test]
625 fn three_penalty_seeds_include_general_spde_manifold_points() {
626 let cfg = SeedConfig::default();
627 let heur = [2.0, 10.0, 3.0];
628 let seeds = generate_rho_candidates(3, Some(&heur), &cfg).expect("ordered seed bounds");
629 let has_non_nu2 = seeds.iter().any(|s| {
630 s.len() == 3 && ((2.0 * s[1] - s[0] - s[2]) - 4.0_f64.ln()).abs() > 1e-3
633 });
634 assert!(has_non_nu2);
635 }
636
637 #[test]
638 fn three_penalty_seeds_include_first_order_fallbackwith_rho2_floor() {
639 let cfg = SeedConfig {
640 bounds: (-12.0, 12.0),
641 ..SeedConfig::default()
642 };
643 let seeds = generate_rho_candidates(3, None, &cfg).expect("ordered seed bounds");
644 let has_floor = seeds
645 .iter()
646 .any(|s| s.len() == 3 && (s[2] - (-12.0)).abs() < 1e-12);
647 assert!(has_floor);
648 }
649
650 #[test]
651 fn auxiliary_trailing_dims_pinned_to_initial_values() {
652 let cfg = SeedConfig {
656 num_auxiliary_trailing: 2,
657 risk_profile: SeedRiskProfile::GeneralizedLinear,
658 ..SeedConfig::default()
659 };
660 let heur = [0.0, 10.0_f64.ln(), 0.0, 0.0]; let seeds = generate_rho_candidates(4, Some(&heur), &cfg).expect("ordered seed bounds");
662 assert!(!seeds.is_empty());
663 for (idx, seed) in seeds.iter().enumerate() {
665 assert_eq!(seed.len(), 4);
666 assert!(
667 (seed[2] - 0.0).abs() < 1e-12 && (seed[3] - 0.0).abs() < 1e-12,
668 "seed {} has auxiliary dims [{}, {}], expected [0, 0]",
669 idx,
670 seed[2],
671 seed[3],
672 );
673 }
674 let has_nonzero_smoothing = seeds
676 .iter()
677 .any(|s| s[0].abs() > 1e-12 || s[1].abs() > 1e-12);
678 assert!(has_nonzero_smoothing);
679 }
680
681 #[test]
682 fn auxiliary_dims_dedup_collapses_identical_seeds() {
683 let cfg = SeedConfig {
686 num_auxiliary_trailing: 1,
687 max_seeds: 32,
688 risk_profile: SeedRiskProfile::GeneralizedLinear,
689 ..SeedConfig::default()
690 };
691 let seeds_with_aux = generate_rho_candidates(3, None, &cfg).expect("ordered seed bounds");
692 let cfg_no_aux = SeedConfig {
693 num_auxiliary_trailing: 0,
694 max_seeds: 32,
695 risk_profile: SeedRiskProfile::GeneralizedLinear,
696 ..SeedConfig::default()
697 };
698 let seeds_without_aux = generate_rho_candidates(3, None, &cfg_no_aux).expect("ordered seed bounds");
699 assert!(seeds_with_aux.len() <= seeds_without_aux.len());
701 }
702}