ipfrs_semantic/
embedding_composer.rs

1//! Embedding Composer — compose multiple embeddings into a single representation
2//! using various late-fusion strategies.
3//!
4//! Supports Concatenate, Average, WeightedAverage, MaxPooling, and HadamardProduct
5//! fusion strategies for building multi-modal or multi-view late-fusion pipelines.
6
7use std::collections::HashMap;
8
9// ---------------------------------------------------------------------------
10// CompositionStrategy
11// ---------------------------------------------------------------------------
12
13/// The fusion strategy to apply when composing multiple embeddings.
14#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
15pub enum CompositionStrategy {
16    /// Concatenate all vectors in order.  Output dim = sum of all input dims.
17    Concatenate,
18    /// Element-wise arithmetic mean.  All inputs must share the same dimension.
19    Average,
20    /// Weighted element-wise average.  All inputs must share the same dimension.
21    /// Each [`EmbeddingInput::weight`] is used; weights are normalised to sum to 1.
22    WeightedAverage,
23    /// Element-wise maximum.  All inputs must share the same dimension.
24    MaxPooling,
25    /// Element-wise product (Hadamard product).  All inputs must share the same dimension.
26    HadamardProduct,
27}
28
29// ---------------------------------------------------------------------------
30// EmbeddingInput
31// ---------------------------------------------------------------------------
32
33/// A single embedding source together with its metadata.
34#[derive(Debug, Clone)]
35pub struct EmbeddingInput {
36    /// Opaque identifier for the originating source (e.g. modality id, model id).
37    pub source_id: u64,
38    /// The embedding vector.
39    pub vector: Vec<f32>,
40    /// Weight used by [`CompositionStrategy::WeightedAverage`].
41    /// Ignored by other strategies.
42    pub weight: f64,
43}
44
45impl EmbeddingInput {
46    /// Convenience constructor.
47    pub fn new(source_id: u64, vector: Vec<f32>, weight: f64) -> Self {
48        Self {
49            source_id,
50            vector,
51            weight,
52        }
53    }
54}
55
56// ---------------------------------------------------------------------------
57// CompositionResult
58// ---------------------------------------------------------------------------
59
60/// The outcome of a composition operation.
61#[derive(Debug, Clone)]
62pub struct CompositionResult {
63    /// The composed embedding vector.
64    pub composed: Vec<f32>,
65    /// The strategy that produced this result.
66    pub strategy: CompositionStrategy,
67    /// Number of input embeddings that were fused.
68    pub input_count: usize,
69    /// Dimensionality of the composed vector.
70    pub output_dim: usize,
71}
72
73impl CompositionResult {
74    /// Compute the L2 (Euclidean) norm of the composed vector.
75    pub fn l2_norm(&self) -> f32 {
76        self.composed.iter().map(|v| v * v).sum::<f32>().sqrt()
77    }
78
79    /// Return an L2-normalised copy of the composed vector.
80    /// If the norm is zero (or very close to zero), returns a zero vector of the same length.
81    pub fn normalize(&self) -> Vec<f32> {
82        let norm = self.l2_norm();
83        if norm == 0.0 {
84            return vec![0.0_f32; self.composed.len()];
85        }
86        self.composed.iter().map(|v| v / norm).collect()
87    }
88}
89
90// ---------------------------------------------------------------------------
91// ComposerStats
92// ---------------------------------------------------------------------------
93
94/// Accumulated statistics for an [`EmbeddingComposer`] instance.
95#[derive(Debug, Clone, Default)]
96pub struct ComposerStats {
97    /// Total number of successful composition operations.
98    pub total_composed: u64,
99    /// Per-strategy success counts.
100    pub strategy_counts: HashMap<CompositionStrategy, u64>,
101}
102
103impl ComposerStats {
104    /// Returns the strategy that has been used most often, or `None` if no
105    /// compositions have been performed yet.
106    pub fn most_used_strategy(&self) -> Option<CompositionStrategy> {
107        self.strategy_counts
108            .iter()
109            .max_by_key(|(_, count)| *count)
110            .map(|(strategy, _)| *strategy)
111    }
112}
113
114// ---------------------------------------------------------------------------
115// EmbeddingComposer
116// ---------------------------------------------------------------------------
117
118/// Composes multiple embeddings into a single representation using a chosen
119/// [`CompositionStrategy`].
120///
121/// # Example
122///
123/// ```rust
124/// use ipfrs_semantic::embedding_composer::{
125///     EmbeddingComposer, EmbeddingInput, CompositionStrategy,
126/// };
127///
128/// let mut composer = EmbeddingComposer::new();
129///
130/// let inputs = vec![
131///     EmbeddingInput::new(1, vec![1.0, 2.0], 1.0),
132///     EmbeddingInput::new(2, vec![3.0, 4.0], 1.0),
133/// ];
134///
135/// let result = composer
136///     .compose(&inputs, CompositionStrategy::Concatenate)
137///     .unwrap();
138/// assert_eq!(result.composed, vec![1.0, 2.0, 3.0, 4.0]);
139/// ```
140#[derive(Debug, Default)]
141pub struct EmbeddingComposer {
142    /// Accumulated statistics across all compose calls.
143    pub stats: ComposerStats,
144}
145
146impl EmbeddingComposer {
147    /// Create a new composer with zeroed statistics.
148    pub fn new() -> Self {
149        Self {
150            stats: ComposerStats::default(),
151        }
152    }
153
154    /// Compose `inputs` using `strategy` and return the result.
155    ///
156    /// # Errors
157    ///
158    /// - Returns `Err("no inputs")` when `inputs` is empty.
159    /// - Returns `Err("dimension mismatch")` when a strategy that requires
160    ///   equal dimensions receives inputs of differing dimensionality.
161    pub fn compose(
162        &mut self,
163        inputs: &[EmbeddingInput],
164        strategy: CompositionStrategy,
165    ) -> Result<CompositionResult, String> {
166        if inputs.is_empty() {
167            return Err("no inputs".to_string());
168        }
169
170        let composed = match strategy {
171            CompositionStrategy::Concatenate => Self::apply_concatenate(inputs),
172            CompositionStrategy::Average => Self::apply_average(inputs)?,
173            CompositionStrategy::WeightedAverage => Self::apply_weighted_average(inputs)?,
174            CompositionStrategy::MaxPooling => Self::apply_max_pooling(inputs)?,
175            CompositionStrategy::HadamardProduct => Self::apply_hadamard(inputs)?,
176        };
177
178        let output_dim = composed.len();
179        let input_count = inputs.len();
180
181        // Update stats.
182        self.stats.total_composed += 1;
183        *self.stats.strategy_counts.entry(strategy).or_insert(0) += 1;
184
185        Ok(CompositionResult {
186            composed,
187            strategy,
188            input_count,
189            output_dim,
190        })
191    }
192
193    /// Compose multiple independent batches in a single call.
194    ///
195    /// Each element of `batches` is a `(inputs, strategy)` pair.  Results are
196    /// returned in the same order as the input batches.
197    pub fn batch_compose(
198        &mut self,
199        batches: Vec<(Vec<EmbeddingInput>, CompositionStrategy)>,
200    ) -> Vec<Result<CompositionResult, String>> {
201        batches
202            .into_iter()
203            .map(|(inputs, strategy)| self.compose(&inputs, strategy))
204            .collect()
205    }
206
207    /// Return a reference to the accumulated statistics.
208    pub fn stats(&self) -> &ComposerStats {
209        &self.stats
210    }
211
212    // -----------------------------------------------------------------------
213    // Private strategy implementations
214    // -----------------------------------------------------------------------
215
216    fn apply_concatenate(inputs: &[EmbeddingInput]) -> Vec<f32> {
217        inputs
218            .iter()
219            .flat_map(|inp| inp.vector.iter().copied())
220            .collect()
221    }
222
223    fn check_same_dim(inputs: &[EmbeddingInput]) -> Result<usize, String> {
224        let dim = inputs[0].vector.len();
225        for inp in inputs.iter().skip(1) {
226            if inp.vector.len() != dim {
227                return Err("dimension mismatch".to_string());
228            }
229        }
230        Ok(dim)
231    }
232
233    fn apply_average(inputs: &[EmbeddingInput]) -> Result<Vec<f32>, String> {
234        let dim = Self::check_same_dim(inputs)?;
235        let n = inputs.len() as f32;
236        let mut result = vec![0.0_f32; dim];
237        for inp in inputs {
238            for (r, v) in result.iter_mut().zip(inp.vector.iter()) {
239                *r += v;
240            }
241        }
242        for r in &mut result {
243            *r /= n;
244        }
245        Ok(result)
246    }
247
248    fn apply_weighted_average(inputs: &[EmbeddingInput]) -> Result<Vec<f32>, String> {
249        let dim = Self::check_same_dim(inputs)?;
250
251        // Normalise weights: if all weights sum to zero, treat all as equal.
252        let weight_sum: f64 = inputs.iter().map(|inp| inp.weight).sum();
253        let weights: Vec<f64> = if weight_sum == 0.0 {
254            let equal = 1.0 / inputs.len() as f64;
255            vec![equal; inputs.len()]
256        } else {
257            inputs.iter().map(|inp| inp.weight / weight_sum).collect()
258        };
259
260        let mut result = vec![0.0_f32; dim];
261        for (inp, w) in inputs.iter().zip(weights.iter()) {
262            let wf = *w as f32;
263            for (r, v) in result.iter_mut().zip(inp.vector.iter()) {
264                *r += wf * v;
265            }
266        }
267        Ok(result)
268    }
269
270    fn apply_max_pooling(inputs: &[EmbeddingInput]) -> Result<Vec<f32>, String> {
271        let dim = Self::check_same_dim(inputs)?;
272        let mut result = inputs[0].vector.clone();
273        for inp in inputs.iter().skip(1) {
274            for (r, v) in result.iter_mut().zip(inp.vector.iter()) {
275                if *v > *r {
276                    *r = *v;
277                }
278            }
279        }
280        // Suppress unused variable warning from `dim` when result is taken from clone.
281        let _ = dim;
282        Ok(result)
283    }
284
285    fn apply_hadamard(inputs: &[EmbeddingInput]) -> Result<Vec<f32>, String> {
286        let dim = Self::check_same_dim(inputs)?;
287        let mut result = inputs[0].vector.clone();
288        for inp in inputs.iter().skip(1) {
289            for (r, v) in result.iter_mut().zip(inp.vector.iter()) {
290                *r *= v;
291            }
292        }
293        let _ = dim;
294        Ok(result)
295    }
296}
297
298// ---------------------------------------------------------------------------
299// Tests
300// ---------------------------------------------------------------------------
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305
306    fn make_input(source_id: u64, vector: Vec<f32>, weight: f64) -> EmbeddingInput {
307        EmbeddingInput::new(source_id, vector, weight)
308    }
309
310    // -------
311    // Concatenate
312    // -------
313
314    #[test]
315    fn test_concatenate_joins_vectors() {
316        let mut c = EmbeddingComposer::new();
317        let inputs = vec![
318            make_input(1, vec![1.0, 2.0], 1.0),
319            make_input(2, vec![3.0, 4.0], 1.0),
320        ];
321        let res = c
322            .compose(&inputs, CompositionStrategy::Concatenate)
323            .expect("test: compose concatenate failed");
324        assert_eq!(res.composed, vec![1.0, 2.0, 3.0, 4.0]);
325    }
326
327    #[test]
328    fn test_concatenate_three_inputs() {
329        let mut c = EmbeddingComposer::new();
330        let inputs = vec![
331            make_input(1, vec![1.0], 1.0),
332            make_input(2, vec![2.0], 1.0),
333            make_input(3, vec![3.0], 1.0),
334        ];
335        let res = c
336            .compose(&inputs, CompositionStrategy::Concatenate)
337            .expect("test: compose concatenate three inputs failed");
338        assert_eq!(res.composed, vec![1.0, 2.0, 3.0]);
339    }
340
341    #[test]
342    fn test_concatenate_output_dim_is_sum() {
343        let mut c = EmbeddingComposer::new();
344        let inputs = vec![
345            make_input(1, vec![0.0; 3], 1.0),
346            make_input(2, vec![0.0; 5], 1.0),
347        ];
348        let res = c
349            .compose(&inputs, CompositionStrategy::Concatenate)
350            .expect("test: compose concatenate output dim failed");
351        assert_eq!(res.output_dim, 8);
352    }
353
354    #[test]
355    fn test_concatenate_different_dims_allowed() {
356        let mut c = EmbeddingComposer::new();
357        let inputs = vec![
358            make_input(1, vec![1.0, 2.0, 3.0], 1.0),
359            make_input(2, vec![4.0], 1.0),
360        ];
361        let res = c
362            .compose(&inputs, CompositionStrategy::Concatenate)
363            .expect("test: compose concatenate different dims failed");
364        assert_eq!(res.composed, vec![1.0, 2.0, 3.0, 4.0]);
365    }
366
367    // -------
368    // Average
369    // -------
370
371    #[test]
372    fn test_average_is_element_mean() {
373        let mut c = EmbeddingComposer::new();
374        let inputs = vec![
375            make_input(1, vec![0.0, 4.0], 1.0),
376            make_input(2, vec![2.0, 0.0], 1.0),
377        ];
378        let res = c
379            .compose(&inputs, CompositionStrategy::Average)
380            .expect("test: compose average failed");
381        assert!((res.composed[0] - 1.0).abs() < 1e-6);
382        assert!((res.composed[1] - 2.0).abs() < 1e-6);
383    }
384
385    #[test]
386    fn test_average_three_inputs() {
387        let mut c = EmbeddingComposer::new();
388        let inputs = vec![
389            make_input(1, vec![3.0], 1.0),
390            make_input(2, vec![6.0], 1.0),
391            make_input(3, vec![9.0], 1.0),
392        ];
393        let res = c
394            .compose(&inputs, CompositionStrategy::Average)
395            .expect("test: compose average three inputs failed");
396        assert!((res.composed[0] - 6.0).abs() < 1e-6);
397    }
398
399    #[test]
400    fn test_average_output_dim_same_as_input() {
401        let mut c = EmbeddingComposer::new();
402        let inputs = vec![
403            make_input(1, vec![0.0; 4], 1.0),
404            make_input(2, vec![0.0; 4], 1.0),
405        ];
406        let res = c
407            .compose(&inputs, CompositionStrategy::Average)
408            .expect("test: compose average output dim failed");
409        assert_eq!(res.output_dim, 4);
410    }
411
412    // -------
413    // WeightedAverage
414    // -------
415
416    #[test]
417    fn test_weighted_average_applies_weights() {
418        let mut c = EmbeddingComposer::new();
419        // weight 1 (0.25) and weight 3 (0.75)
420        let inputs = vec![make_input(1, vec![0.0], 1.0), make_input(2, vec![4.0], 3.0)];
421        let res = c
422            .compose(&inputs, CompositionStrategy::WeightedAverage)
423            .expect("test: compose weighted average failed");
424        // Expected: 0*0.25 + 4*0.75 = 3.0
425        assert!((res.composed[0] - 3.0).abs() < 1e-5);
426    }
427
428    #[test]
429    fn test_weighted_average_zero_weight_sum_is_equal() {
430        let mut c = EmbeddingComposer::new();
431        let inputs = vec![make_input(1, vec![0.0], 0.0), make_input(2, vec![4.0], 0.0)];
432        let res = c
433            .compose(&inputs, CompositionStrategy::WeightedAverage)
434            .expect("test: compose weighted average zero weights failed");
435        // Equal weights → mean
436        assert!((res.composed[0] - 2.0).abs() < 1e-5);
437    }
438
439    // -------
440    // MaxPooling
441    // -------
442
443    #[test]
444    fn test_max_pooling_takes_element_max() {
445        let mut c = EmbeddingComposer::new();
446        let inputs = vec![
447            make_input(1, vec![1.0, 5.0, 2.0], 1.0),
448            make_input(2, vec![3.0, 2.0, 7.0], 1.0),
449        ];
450        let res = c
451            .compose(&inputs, CompositionStrategy::MaxPooling)
452            .expect("test: compose max pooling failed");
453        assert_eq!(res.composed, vec![3.0, 5.0, 7.0]);
454    }
455
456    #[test]
457    fn test_max_pooling_three_inputs() {
458        let mut c = EmbeddingComposer::new();
459        let inputs = vec![
460            make_input(1, vec![1.0, 9.0], 1.0),
461            make_input(2, vec![5.0, 2.0], 1.0),
462            make_input(3, vec![3.0, 7.0], 1.0),
463        ];
464        let res = c
465            .compose(&inputs, CompositionStrategy::MaxPooling)
466            .expect("test: compose max pooling three inputs failed");
467        assert_eq!(res.composed, vec![5.0, 9.0]);
468    }
469
470    // -------
471    // HadamardProduct
472    // -------
473
474    #[test]
475    fn test_hadamard_product_multiplies() {
476        let mut c = EmbeddingComposer::new();
477        let inputs = vec![
478            make_input(1, vec![2.0, 3.0], 1.0),
479            make_input(2, vec![4.0, 5.0], 1.0),
480        ];
481        let res = c
482            .compose(&inputs, CompositionStrategy::HadamardProduct)
483            .expect("test: compose hadamard product failed");
484        assert_eq!(res.composed, vec![8.0, 15.0]);
485    }
486
487    #[test]
488    fn test_hadamard_product_three_inputs() {
489        let mut c = EmbeddingComposer::new();
490        let inputs = vec![
491            make_input(1, vec![2.0], 1.0),
492            make_input(2, vec![3.0], 1.0),
493            make_input(3, vec![4.0], 1.0),
494        ];
495        let res = c
496            .compose(&inputs, CompositionStrategy::HadamardProduct)
497            .expect("test: compose hadamard product three inputs failed");
498        assert_eq!(res.composed, vec![24.0]);
499    }
500
501    // -------
502    // Error cases
503    // -------
504
505    #[test]
506    fn test_empty_inputs_returns_error() {
507        let mut c = EmbeddingComposer::new();
508        let err = c
509            .compose(&[], CompositionStrategy::Average)
510            .expect_err("test: expected error for empty inputs");
511        assert_eq!(err, "no inputs");
512    }
513
514    #[test]
515    fn test_dimension_mismatch_average() {
516        let mut c = EmbeddingComposer::new();
517        let inputs = vec![
518            make_input(1, vec![1.0, 2.0], 1.0),
519            make_input(2, vec![3.0], 1.0),
520        ];
521        let err = c
522            .compose(&inputs, CompositionStrategy::Average)
523            .expect_err("test: expected dimension mismatch error for average");
524        assert_eq!(err, "dimension mismatch");
525    }
526
527    #[test]
528    fn test_dimension_mismatch_max_pooling() {
529        let mut c = EmbeddingComposer::new();
530        let inputs = vec![
531            make_input(1, vec![1.0, 2.0], 1.0),
532            make_input(2, vec![3.0, 4.0, 5.0], 1.0),
533        ];
534        let err = c
535            .compose(&inputs, CompositionStrategy::MaxPooling)
536            .expect_err("test: expected dimension mismatch error for max pooling");
537        assert_eq!(err, "dimension mismatch");
538    }
539
540    #[test]
541    fn test_dimension_mismatch_hadamard() {
542        let mut c = EmbeddingComposer::new();
543        let inputs = vec![
544            make_input(1, vec![1.0], 1.0),
545            make_input(2, vec![1.0, 2.0], 1.0),
546        ];
547        let err = c
548            .compose(&inputs, CompositionStrategy::HadamardProduct)
549            .expect_err("test: expected dimension mismatch error for hadamard");
550        assert_eq!(err, "dimension mismatch");
551    }
552
553    // -------
554    // l2_norm and normalize
555    // -------
556
557    #[test]
558    fn test_l2_norm_correct() {
559        let mut c = EmbeddingComposer::new();
560        let inputs = vec![make_input(1, vec![3.0, 4.0], 1.0)];
561        let res = c
562            .compose(&inputs, CompositionStrategy::Concatenate)
563            .expect("test: compose for l2 norm failed");
564        assert!((res.l2_norm() - 5.0).abs() < 1e-6);
565    }
566
567    #[test]
568    fn test_normalize_unit_length() {
569        let mut c = EmbeddingComposer::new();
570        let inputs = vec![make_input(1, vec![0.0, 3.0, 4.0], 1.0)];
571        let res = c
572            .compose(&inputs, CompositionStrategy::Concatenate)
573            .expect("test: compose for normalize failed");
574        let norm_vec = res.normalize();
575        let sq_sum: f32 = norm_vec.iter().map(|v| v * v).sum();
576        assert!((sq_sum - 1.0).abs() < 1e-6);
577    }
578
579    #[test]
580    fn test_normalize_zero_vector_returns_zeros() {
581        let mut c = EmbeddingComposer::new();
582        let inputs = vec![make_input(1, vec![0.0, 0.0], 1.0)];
583        let res = c
584            .compose(&inputs, CompositionStrategy::Concatenate)
585            .expect("test: compose for normalize zero vector failed");
586        assert_eq!(res.normalize(), vec![0.0, 0.0]);
587    }
588
589    // -------
590    // batch_compose
591    // -------
592
593    #[test]
594    fn test_batch_compose_returns_correct_count() {
595        let mut c = EmbeddingComposer::new();
596        let batches = vec![
597            (
598                vec![
599                    make_input(1, vec![1.0, 2.0], 1.0),
600                    make_input(2, vec![3.0, 4.0], 1.0),
601                ],
602                CompositionStrategy::Concatenate,
603            ),
604            (
605                vec![make_input(3, vec![1.0], 1.0), make_input(4, vec![1.0], 1.0)],
606                CompositionStrategy::Average,
607            ),
608        ];
609        let results = c.batch_compose(batches);
610        assert_eq!(results.len(), 2);
611        assert!(results[0].is_ok());
612        assert!(results[1].is_ok());
613    }
614
615    #[test]
616    fn test_batch_compose_with_error() {
617        let mut c = EmbeddingComposer::new();
618        let batches = vec![
619            // Valid
620            (
621                vec![make_input(1, vec![1.0], 1.0), make_input(2, vec![2.0], 1.0)],
622                CompositionStrategy::Average,
623            ),
624            // Invalid: empty
625            (vec![], CompositionStrategy::MaxPooling),
626        ];
627        let results = c.batch_compose(batches);
628        assert!(results[0].is_ok());
629        assert!(results[1].is_err());
630    }
631
632    // -------
633    // Stats
634    // -------
635
636    #[test]
637    fn test_stats_total_composed() {
638        let mut c = EmbeddingComposer::new();
639        let inputs = vec![make_input(1, vec![1.0], 1.0)];
640        c.compose(&inputs, CompositionStrategy::Concatenate)
641            .expect("test: first compose for stats failed");
642        c.compose(&inputs, CompositionStrategy::Concatenate)
643            .expect("test: second compose for stats failed");
644        assert_eq!(c.stats().total_composed, 2);
645    }
646
647    #[test]
648    fn test_stats_most_used_strategy() {
649        let mut c = EmbeddingComposer::new();
650        let inputs = vec![make_input(1, vec![1.0], 1.0), make_input(2, vec![1.0], 1.0)];
651        c.compose(&inputs, CompositionStrategy::Average)
652            .expect("test: compose average for stats failed");
653        c.compose(&inputs, CompositionStrategy::Average)
654            .expect("test: compose second average for stats failed");
655        c.compose(&inputs, CompositionStrategy::MaxPooling)
656            .expect("test: compose max pooling for stats failed");
657        assert_eq!(
658            c.stats().most_used_strategy(),
659            Some(CompositionStrategy::Average)
660        );
661    }
662
663    #[test]
664    fn test_stats_most_used_strategy_none_when_empty() {
665        let c = EmbeddingComposer::new();
666        assert_eq!(c.stats().most_used_strategy(), None);
667    }
668
669    // -------
670    // input_count field
671    // -------
672
673    #[test]
674    fn test_input_count_field() {
675        let mut c = EmbeddingComposer::new();
676        let inputs = vec![
677            make_input(1, vec![1.0], 1.0),
678            make_input(2, vec![2.0], 1.0),
679            make_input(3, vec![3.0], 1.0),
680        ];
681        let res = c
682            .compose(&inputs, CompositionStrategy::Concatenate)
683            .expect("test: compose for input_count failed");
684        assert_eq!(res.input_count, 3);
685    }
686}
ipfrs_semantic/embedding_composer.rs

ipfrs_semantic/
embedding_composer.rs
