pictor_model/model_merge.rs
1//! Model merging utilities: linear interpolation, SLERP, TIES, and task-vector merging.
2//!
3//! These methods allow creating new model variants by combining weights from
4//! multiple source models — useful for model fusion, continual learning, and
5//! experimental model architectures.
6//!
7//! ## Available Methods
8//!
9//! - **Linear**: simple weighted average `(1-α)*A + α*B`
10//! - **SLERP**: spherical linear interpolation (preserves direction on the unit hypersphere)
11//! - **TIES**: sign-majority election with magnitude trimming (from the TIES-Merging paper)
12//! - **Task Vector**: `base + α*(finetuned - base)`, adds or subtracts fine-tuning direction
13//! - **DARE**: random dropout of task-vector elements with rescaling for sparse merging
14//!
15//! ## Example
16//!
17//! ```rust
18//! use pictor_model::model_merge::{WeightTensor, MergeConfig, MergeMethod, merge_models};
19//!
20//! let base = vec![
21//! WeightTensor::new("embed.weight", vec![1.0, 0.0, 0.0, 1.0], vec![2, 2]),
22//! ];
23//! let other = vec![
24//! WeightTensor::new("embed.weight", vec![0.0, 1.0, 1.0, 0.0], vec![2, 2]),
25//! ];
26//! let config = MergeConfig { method: MergeMethod::Linear, alpha: 0.5, ..Default::default() };
27//! let merged = merge_models(&base, &other, &config).expect("merge failed");
28//! assert_eq!(merged.len(), 1);
29//! ```
30
31use std::collections::HashMap;
32use thiserror::Error;
33
34// ──────────────────────────────────────────────────────────────────
35// Error type
36// ──────────────────────────────────────────────────────────────────
37
38/// Errors that can occur during model merging operations.
39#[derive(Debug, Error)]
40pub enum MergeError {
41 /// Shape mismatch between the two tensors being merged.
42 #[error("shape mismatch for tensor '{name}': {a:?} vs {b:?}")]
43 ShapeMismatch {
44 name: String,
45 a: Vec<usize>,
46 b: Vec<usize>,
47 },
48 /// Tensor has no elements (empty data slice or shape contains zero).
49 #[error("empty tensor: '{0}'")]
50 EmptyTensor(String),
51 /// Alpha coefficient is outside the valid range `[0.0, 1.0]`.
52 #[error("invalid alpha {0}: must be in [0.0, 1.0]")]
53 InvalidAlpha(f32),
54 /// Density is outside the valid range `(0.0, 1.0]`.
55 #[error("invalid density {0}: must be in (0.0, 1.0]")]
56 InvalidDensity(f32),
57 /// SLERP was attempted on a zero-norm vector.
58 #[error("SLERP failed: zero vector")]
59 SierpZeroVector,
60}
61
62// ──────────────────────────────────────────────────────────────────
63// WeightTensor
64// ──────────────────────────────────────────────────────────────────
65
66/// A named weight tensor — a flat `f32` slice with shape metadata.
67///
68/// The `data` field is stored in row-major (C-contiguous) order.
69/// The product of all shape dimensions must equal `data.len()`.
70#[derive(Debug, Clone)]
71pub struct WeightTensor {
72 /// Unique name identifying this tensor within a model checkpoint.
73 pub name: String,
74 /// Raw weight data in row-major order.
75 pub data: Vec<f32>,
76 /// N-dimensional shape (e.g., `[4096, 4096]` for a square linear layer).
77 pub shape: Vec<usize>,
78}
79
80impl WeightTensor {
81 /// Construct a weight tensor from its components.
82 pub fn new(name: impl Into<String>, data: Vec<f32>, shape: Vec<usize>) -> Self {
83 Self {
84 name: name.into(),
85 data,
86 shape,
87 }
88 }
89
90 /// Construct an all-zeros tensor with the given name and shape.
91 pub fn zeros(name: impl Into<String>, shape: Vec<usize>) -> Self {
92 let n = shape.iter().product();
93 Self {
94 name: name.into(),
95 data: vec![0.0f32; n],
96 shape,
97 }
98 }
99
100 /// Number of scalar elements: product of all shape dimensions.
101 pub fn element_count(&self) -> usize {
102 self.shape.iter().product()
103 }
104
105 /// Euclidean (L2) norm: `sqrt(sum(x_i^2))`.
106 pub fn l2_norm(&self) -> f32 {
107 self.data.iter().map(|x| x * x).sum::<f32>().sqrt()
108 }
109
110 /// Cosine similarity with `other`: `dot(a, b) / (|a| * |b|)`.
111 ///
112 /// Returns `Err(MergeError::ShapeMismatch)` when element counts differ,
113 /// `Err(MergeError::EmptyTensor)` when either tensor is empty.
114 pub fn cosine_similarity(&self, other: &WeightTensor) -> Result<f32, MergeError> {
115 let n = self.element_count();
116 if n == 0 {
117 return Err(MergeError::EmptyTensor(self.name.clone()));
118 }
119 if other.element_count() == 0 {
120 return Err(MergeError::EmptyTensor(other.name.clone()));
121 }
122 if n != other.element_count() {
123 return Err(MergeError::ShapeMismatch {
124 name: self.name.clone(),
125 a: self.shape.clone(),
126 b: other.shape.clone(),
127 });
128 }
129
130 let dot: f32 = self
131 .data
132 .iter()
133 .zip(other.data.iter())
134 .map(|(a, b)| a * b)
135 .sum();
136 let norm_a = self.l2_norm();
137 let norm_b = other.l2_norm();
138 let denom = norm_a * norm_b;
139 if denom == 0.0 {
140 // At least one zero vector — cosine similarity is conventionally 0.
141 return Ok(0.0);
142 }
143 Ok(dot / denom)
144 }
145
146 /// Element-wise addition: `self + other`.
147 pub fn add(&self, other: &WeightTensor) -> Result<WeightTensor, MergeError> {
148 check_compatible(self, other)?;
149 let data: Vec<f32> = self
150 .data
151 .iter()
152 .zip(other.data.iter())
153 .map(|(a, b)| a + b)
154 .collect();
155 Ok(WeightTensor::new(
156 self.name.clone(),
157 data,
158 self.shape.clone(),
159 ))
160 }
161
162 /// Element-wise subtraction: `self - other`.
163 pub fn sub(&self, other: &WeightTensor) -> Result<WeightTensor, MergeError> {
164 check_compatible(self, other)?;
165 let data: Vec<f32> = self
166 .data
167 .iter()
168 .zip(other.data.iter())
169 .map(|(a, b)| a - b)
170 .collect();
171 Ok(WeightTensor::new(
172 self.name.clone(),
173 data,
174 self.shape.clone(),
175 ))
176 }
177
178 /// Scalar multiplication: `self * alpha`.
179 pub fn scale(&self, alpha: f32) -> WeightTensor {
180 let data: Vec<f32> = self.data.iter().map(|x| x * alpha).collect();
181 WeightTensor::new(self.name.clone(), data, self.shape.clone())
182 }
183
184 /// Linear interpolation: `(1 - t)*self + t*other`.
185 ///
186 /// `t` is not validated here — use [`merge_tensors`] for validated entry points.
187 pub fn lerp(&self, other: &WeightTensor, t: f32) -> Result<WeightTensor, MergeError> {
188 check_compatible(self, other)?;
189 let data = linear_merge(&self.data, &other.data, t);
190 Ok(WeightTensor::new(
191 self.name.clone(),
192 data,
193 self.shape.clone(),
194 ))
195 }
196}
197
198// ──────────────────────────────────────────────────────────────────
199// MergeMethod / MergeConfig
200// ──────────────────────────────────────────────────────────────────
201
202/// Available merge strategies.
203#[derive(Debug, Clone, PartialEq)]
204pub enum MergeMethod {
205 /// Simple weighted average: `result = (1-α)*A + α*B`.
206 Linear,
207 /// Spherical linear interpolation — preserves direction on the unit hypersphere.
208 Slerp,
209 /// TIES-Merging: sign-majority election with magnitude trimming.
210 Ties,
211 /// Task-vector merging: `base + α*(finetuned - base)`.
212 ///
213 /// When `α=1.0` this is identical to returning `finetuned`;
214 /// when `α=0.0` it is identical to returning `base`.
215 TaskVector,
216 /// DARE (Drop And REscale): deterministically drop task-vector elements, then rescale.
217 Dare {
218 /// Seed for the deterministic LCG random number generator.
219 seed: u64,
220 /// Fraction of task-vector elements to zero out (0.0 = keep all, 1.0 = drop all).
221 dropout_rate: f32,
222 },
223}
224
225/// Configuration for a single merge pass.
226#[derive(Debug, Clone)]
227pub struct MergeConfig {
228 /// Which merging algorithm to use.
229 pub method: MergeMethod,
230 /// Interpolation coefficient: 0.0 → model A only, 1.0 → model B only.
231 pub alpha: f32,
232 /// Normalize each weight tensor to unit L2 norm before merging.
233 pub normalize: bool,
234 /// For TIES: fraction of weights to retain by magnitude (0.0–1.0].
235 pub density: f32,
236}
237
238impl Default for MergeConfig {
239 fn default() -> Self {
240 Self {
241 method: MergeMethod::Linear,
242 alpha: 0.5,
243 normalize: false,
244 density: 0.5,
245 }
246 }
247}
248
249// ──────────────────────────────────────────────────────────────────
250// MergeStats
251// ──────────────────────────────────────────────────────────────────
252
253/// Statistics collected during a model merge operation.
254#[derive(Debug, Clone)]
255pub struct MergeStats {
256 /// Number of tensors that were present in both models and were merged.
257 pub tensors_merged: usize,
258 /// Number of tensors that existed only in `base` and were copied unchanged.
259 pub tensors_copied: usize,
260 /// Total number of scalar parameters in the output model.
261 pub total_params: usize,
262 /// Average cosine similarity between corresponding base/other tensors.
263 pub mean_cosine_similarity: f32,
264 /// The merge method that was used.
265 pub method: MergeMethod,
266}
267
268impl MergeStats {
269 /// Human-readable one-line summary of the merge statistics.
270 pub fn summary(&self) -> String {
271 format!(
272 "method={:?} merged={} copied={} total_params={} mean_cosine_sim={:.4}",
273 self.method,
274 self.tensors_merged,
275 self.tensors_copied,
276 self.total_params,
277 self.mean_cosine_similarity,
278 )
279 }
280}
281
282// ──────────────────────────────────────────────────────────────────
283// Low-level primitive functions
284// ──────────────────────────────────────────────────────────────────
285
286/// Linear interpolation element-wise: `result[i] = (1-α)*a[i] + α*b[i]`.
287///
288/// `a` and `b` must be the same length; if they differ the shorter slice
289/// determines the output length (extra elements from the longer slice are dropped).
290pub fn linear_merge(a: &[f32], b: &[f32], alpha: f32) -> Vec<f32> {
291 let one_minus_alpha = 1.0 - alpha;
292 a.iter()
293 .zip(b.iter())
294 .map(|(ai, bi)| one_minus_alpha * ai + alpha * bi)
295 .collect()
296}
297
298/// Spherical linear interpolation (SLERP) between two real-valued vectors.
299///
300/// Both vectors are first normalized to unit length. If either has zero norm
301/// or if they are nearly parallel (`cos_theta > 0.9995`), the function falls back
302/// to ordinary linear interpolation to avoid numerical instability.
303///
304/// ## Formula
305///
306/// `result = sin((1-t)*θ)/sin(θ) * a + sin(t*θ)/sin(θ) * b`
307///
308/// where `θ = acos(dot(a_norm, b_norm))`.
309pub fn slerp(a: &[f32], b: &[f32], t: f32) -> Vec<f32> {
310 let n = a.len().min(b.len());
311 if n == 0 {
312 return Vec::new();
313 }
314
315 // Compute norms
316 let norm_a = a.iter().map(|x| x * x).sum::<f32>().sqrt();
317 let norm_b = b.iter().map(|x| x * x).sum::<f32>().sqrt();
318
319 // Fall back to linear when either vector is a zero vector
320 if norm_a < f32::EPSILON || norm_b < f32::EPSILON {
321 return linear_merge(a, b, t);
322 }
323
324 // Dot product of normalized vectors
325 let cos_theta: f32 = a[..n]
326 .iter()
327 .zip(b[..n].iter())
328 .map(|(ai, bi)| (ai / norm_a) * (bi / norm_b))
329 .sum::<f32>()
330 .clamp(-1.0, 1.0);
331
332 // Nearly parallel: fall back to linear
333 if cos_theta > 0.9995 {
334 return linear_merge(a, b, t);
335 }
336
337 let theta = cos_theta.acos();
338 let sin_theta = theta.sin();
339
340 // Safety: sin_theta should be > 0 here since |cos_theta| < 0.9995
341 if sin_theta.abs() < f32::EPSILON {
342 return linear_merge(a, b, t);
343 }
344
345 let coeff_a = ((1.0 - t) * theta).sin() / sin_theta;
346 let coeff_b = (t * theta).sin() / sin_theta;
347
348 a[..n]
349 .iter()
350 .zip(b[..n].iter())
351 .map(|(ai, bi)| coeff_a * ai + coeff_b * bi)
352 .collect()
353}
354
355/// TIES-Merging: magnitude-based trimming followed by sign-majority election.
356///
357/// ## Algorithm
358///
359/// 1. Compute delta vectors: `δa = a - mean(a)`, `δb = b - mean(b)` — but for
360/// checkpoint merging we treat `a` and `b` directly as task vectors.
361/// 2. **Trim**: for each of `a` and `b` independently, zero out the bottom
362/// `(1 - density)` fraction of weights by absolute magnitude.
363/// 3. **Elect**: for each position, keep the delta whose magnitude is larger;
364/// if they agree in sign use their average scaled by `alpha`, otherwise use
365/// the one with larger magnitude scaled by `alpha`.
366/// 4. **Result**: `0.5*(a_trimmed + b_trimmed)` weighted by alpha on the
367/// agreement region; positions where signs conflict take the dominant sign.
368///
369/// `density` should be in `(0.0, 1.0]`; 1.0 keeps all weights (no trimming).
370pub fn ties_merge(a: &[f32], b: &[f32], alpha: f32, density: f32) -> Vec<f32> {
371 let n = a.len().min(b.len());
372 if n == 0 {
373 return Vec::new();
374 }
375
376 // --- Trim step ---
377 let trimmed_a = trim_by_magnitude(a, density);
378 let trimmed_b = trim_by_magnitude(b, density);
379
380 // --- Sign-majority + magnitude-dominant election ---
381 trimmed_a
382 .iter()
383 .zip(trimmed_b.iter())
384 .map(|(va, vb)| {
385 let sign_a = va.signum(); // -1.0, 0.0, or 1.0
386 let sign_b = vb.signum();
387 let abs_a = va.abs();
388 let abs_b = vb.abs();
389
390 if sign_a == sign_b {
391 // Agree in sign: use weighted average (biased by alpha)
392 (1.0 - alpha) * va + alpha * vb
393 } else if abs_a >= abs_b {
394 // Disagree: dominant sign wins (scale by alpha for blending)
395 va * (1.0 - alpha)
396 } else {
397 vb * alpha
398 }
399 })
400 .collect()
401}
402
403/// Task-vector merge: `result[i] = base[i] + alpha * (finetuned[i] - base[i])`.
404///
405/// - `α = 0.0` → returns `base` unchanged
406/// - `α = 1.0` → returns `finetuned` unchanged
407/// - Intermediate values apply a fraction of the fine-tuning direction
408pub fn task_vector_merge(base: &[f32], finetuned: &[f32], alpha: f32) -> Vec<f32> {
409 base.iter()
410 .zip(finetuned.iter())
411 .map(|(b, f)| b + alpha * (f - b))
412 .collect()
413}
414
415/// DARE merge: deterministic drop-and-rescale of task-vector elements.
416///
417/// ## Algorithm
418///
419/// 1. Compute task vector `δ[i] = finetuned[i] - base[i]`.
420/// 2. For each element, use an LCG RNG (seeded with `seed`) to decide whether to
421/// zero it out (with probability `dropout_rate`).
422/// 3. Rescale surviving elements by `1 / (1 - dropout_rate)` to preserve expected
423/// magnitude (analogous to dropout rescaling in neural networks).
424/// 4. Return `base[i] + alpha * δ_sparse[i]`.
425///
426/// The LCG is deterministic: same `seed` always produces the same sparsity mask.
427pub fn dare_merge(
428 base: &[f32],
429 finetuned: &[f32],
430 alpha: f32,
431 dropout_rate: f32,
432 seed: u64,
433) -> Vec<f32> {
434 let mut state = seed;
435 let rescale = if dropout_rate < 1.0 {
436 1.0 / (1.0 - dropout_rate)
437 } else {
438 0.0
439 };
440
441 base.iter()
442 .zip(finetuned.iter())
443 .map(|(b, f)| {
444 let rand_val = lcg_next(&mut state);
445 let delta = f - b;
446 let sparse_delta = if rand_val < dropout_rate {
447 0.0
448 } else {
449 delta * rescale
450 };
451 b + alpha * sparse_delta
452 })
453 .collect()
454}
455
456// ──────────────────────────────────────────────────────────────────
457// High-level tensor-level API
458// ──────────────────────────────────────────────────────────────────
459
460/// Merge two [`WeightTensor`]s using the configured method.
461///
462/// Both tensors must have the same element count. If `config.normalize` is set,
463/// each tensor is scaled to unit L2 norm before merging.
464pub fn merge_tensors(
465 base: &WeightTensor,
466 other: &WeightTensor,
467 config: &MergeConfig,
468) -> Result<WeightTensor, MergeError> {
469 validate_config(config)?;
470 check_compatible(base, other)?;
471 if base.element_count() == 0 {
472 return Err(MergeError::EmptyTensor(base.name.clone()));
473 }
474
475 let (a_data, b_data) = if config.normalize {
476 let norm_a = base.l2_norm();
477 let norm_b = other.l2_norm();
478 let a_norm = if norm_a > f32::EPSILON {
479 base.data.iter().map(|x| x / norm_a).collect()
480 } else {
481 base.data.clone()
482 };
483 let b_norm = if norm_b > f32::EPSILON {
484 other.data.iter().map(|x| x / norm_b).collect()
485 } else {
486 other.data.clone()
487 };
488 (a_norm, b_norm)
489 } else {
490 (base.data.clone(), other.data.clone())
491 };
492
493 let merged_data = apply_merge_method(&a_data, &b_data, config)?;
494 Ok(WeightTensor::new(
495 base.name.clone(),
496 merged_data,
497 base.shape.clone(),
498 ))
499}
500
501/// Merge a full model (collection of named tensors) from two sources.
502///
503/// - Tensors present in **both** models are merged using `config`.
504/// - Tensors present only in `base` are copied unchanged into the output.
505/// - Tensors present only in `other` are silently ignored.
506///
507/// The output preserves the ordering of `base`.
508pub fn merge_models(
509 base: &[WeightTensor],
510 other: &[WeightTensor],
511 config: &MergeConfig,
512) -> Result<Vec<WeightTensor>, MergeError> {
513 let (merged, _stats) = merge_models_with_stats(base, other, config)?;
514 Ok(merged)
515}
516
517/// Merge a full model with statistics collection.
518///
519/// Returns both the merged weight tensors and a [`MergeStats`] summary.
520pub fn merge_models_with_stats(
521 base: &[WeightTensor],
522 other: &[WeightTensor],
523 config: &MergeConfig,
524) -> Result<(Vec<WeightTensor>, MergeStats), MergeError> {
525 validate_config(config)?;
526
527 // Build a name → index lookup for `other`
528 let other_map: HashMap<&str, &WeightTensor> =
529 other.iter().map(|t| (t.name.as_str(), t)).collect();
530
531 let mut result = Vec::with_capacity(base.len());
532 let mut tensors_merged = 0usize;
533 let mut tensors_copied = 0usize;
534 let mut total_params = 0usize;
535 let mut cosine_sum = 0.0f32;
536 let mut cosine_count = 0usize;
537
538 for base_tensor in base {
539 total_params += base_tensor.element_count();
540 if let Some(other_tensor) = other_map.get(base_tensor.name.as_str()) {
541 // Accumulate cosine similarity for stats (best-effort; skip on error)
542 if let Ok(sim) = base_tensor.cosine_similarity(other_tensor) {
543 cosine_sum += sim;
544 cosine_count += 1;
545 }
546
547 let merged_tensor = merge_tensors(base_tensor, other_tensor, config)?;
548 result.push(merged_tensor);
549 tensors_merged += 1;
550 } else {
551 result.push(base_tensor.clone());
552 tensors_copied += 1;
553 }
554 }
555
556 let mean_cosine_similarity = if cosine_count > 0 {
557 cosine_sum / cosine_count as f32
558 } else {
559 0.0
560 };
561
562 let stats = MergeStats {
563 tensors_merged,
564 tensors_copied,
565 total_params,
566 mean_cosine_similarity,
567 method: config.method.clone(),
568 };
569
570 Ok((result, stats))
571}
572
573// ──────────────────────────────────────────────────────────────────
574// Internal helpers
575// ──────────────────────────────────────────────────────────────────
576
577/// Validate that `config.alpha` and `config.density` are in their allowed ranges.
578fn validate_config(config: &MergeConfig) -> Result<(), MergeError> {
579 if !(0.0..=1.0).contains(&config.alpha) {
580 return Err(MergeError::InvalidAlpha(config.alpha));
581 }
582 if config.density <= 0.0 || config.density > 1.0 {
583 return Err(MergeError::InvalidDensity(config.density));
584 }
585 Ok(())
586}
587
588/// Return `Err(ShapeMismatch)` when two tensors have incompatible element counts.
589fn check_compatible(a: &WeightTensor, b: &WeightTensor) -> Result<(), MergeError> {
590 if a.element_count() != b.element_count() {
591 return Err(MergeError::ShapeMismatch {
592 name: a.name.clone(),
593 a: a.shape.clone(),
594 b: b.shape.clone(),
595 });
596 }
597 Ok(())
598}
599
600/// Dispatch to the appropriate primitive merge function based on `config.method`.
601fn apply_merge_method(a: &[f32], b: &[f32], config: &MergeConfig) -> Result<Vec<f32>, MergeError> {
602 match &config.method {
603 MergeMethod::Linear => Ok(linear_merge(a, b, config.alpha)),
604 MergeMethod::Slerp => {
605 // Validate no zero-vector before slerp
606 let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
607 let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
608 if norm_a < f32::EPSILON || norm_b < f32::EPSILON {
609 return Err(MergeError::SierpZeroVector);
610 }
611 Ok(slerp(a, b, config.alpha))
612 }
613 MergeMethod::Ties => Ok(ties_merge(a, b, config.alpha, config.density)),
614 MergeMethod::TaskVector => Ok(task_vector_merge(a, b, config.alpha)),
615 MergeMethod::Dare { seed, dropout_rate } => {
616 Ok(dare_merge(a, b, config.alpha, *dropout_rate, *seed))
617 }
618 }
619}
620
621/// Trim the bottom `(1 - density)` fraction of elements by absolute magnitude.
622///
623/// Elements with magnitude below the computed threshold are zeroed out.
624/// With `density = 1.0` all elements are kept (no-op). With `density = 0.5`
625/// the lower half by magnitude is zeroed.
626fn trim_by_magnitude(data: &[f32], density: f32) -> Vec<f32> {
627 if data.is_empty() {
628 return Vec::new();
629 }
630 if density >= 1.0 {
631 return data.to_vec();
632 }
633
634 // Collect absolute values and sort to find threshold
635 let mut abs_sorted: Vec<f32> = data.iter().map(|x| x.abs()).collect();
636 abs_sorted.sort_by(|x, y| x.partial_cmp(y).unwrap_or(std::cmp::Ordering::Equal));
637
638 // The index below which we trim: trim (1-density) fraction.
639 // Use round() to avoid float-precision edge cases (e.g. 0.4*5 = 1.9999... → 1 without round).
640 let trim_count = ((1.0 - density) * abs_sorted.len() as f32).round() as usize;
641 let threshold = if trim_count < abs_sorted.len() {
642 abs_sorted[trim_count]
643 } else {
644 f32::MAX
645 };
646
647 data.iter()
648 .map(|x| if x.abs() < threshold { 0.0 } else { *x })
649 .collect()
650}
651
652/// Deterministic LCG (linear congruential generator) producing values in `[0.0, 1.0)`.
653///
654/// Uses the Knuth/MMIX parameters:
655/// `state = state * 6364136223846793005 + 1442695040888963407`
656///
657/// The upper 32 bits of the new state are extracted and mapped to `[0, 1)` by
658/// dividing by `2^32` (i.e., `u32::MAX + 1`). This gives a full-range uniform
659/// distribution over the 32-bit space.
660#[inline]
661fn lcg_next(state: &mut u64) -> f32 {
662 *state = state
663 .wrapping_mul(6_364_136_223_846_793_005)
664 .wrapping_add(1_442_695_040_888_963_407);
665 // Extract upper 32 bits and scale to [0, 1)
666 let bits = (*state >> 32) as u32;
667 (bits as f32) / (u32::MAX as f32 + 1.0)
668}
669
670// ──────────────────────────────────────────────────────────────────
671// Unit tests (in-module smoke tests; full suite in tests/)
672// ──────────────────────────────────────────────────────────────────
673
674#[cfg(test)]
675mod tests {
676 use super::*;
677
678 #[test]
679 fn lcg_produces_values_in_unit_interval() {
680 let mut state = 42u64;
681 for _ in 0..1000 {
682 let v = lcg_next(&mut state);
683 assert!((0.0..=1.0).contains(&v), "lcg value {v} out of [0,1]");
684 }
685 }
686
687 #[test]
688 fn trim_by_magnitude_density_one_noop() {
689 let data = vec![0.1, 0.5, -0.3, 0.9, -0.7];
690 let trimmed = trim_by_magnitude(&data, 1.0);
691 assert_eq!(trimmed, data);
692 }
693
694 #[test]
695 fn trim_by_magnitude_zeros_smallest() {
696 // density = 0.6 → keep top 60%, trim bottom 40%
697 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
698 let trimmed = trim_by_magnitude(&data, 0.6);
699 // Bottom 40% of 5 = 2 elements (1.0 and 2.0) should be zeroed
700 assert_eq!(trimmed[0], 0.0, "1.0 should be trimmed");
701 assert_eq!(trimmed[1], 0.0, "2.0 should be trimmed");
702 assert!(trimmed[2] != 0.0, "3.0 should be kept");
703 }
704
705 #[test]
706 fn validate_config_rejects_bad_alpha() {
707 let config = MergeConfig {
708 alpha: 1.5,
709 ..Default::default()
710 };
711 assert!(validate_config(&config).is_err());
712 }
713
714 #[test]
715 fn validate_config_rejects_zero_density() {
716 let config = MergeConfig {
717 density: 0.0,
718 ..Default::default()
719 };
720 assert!(validate_config(&config).is_err());
721 }
722}