Skip to main content

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}