tensor4all_core/
index_ops.rs

1use crate::IndexLike;
2use tinyvec::TinyVec;
3
4const SMALL_CONTRACTION_INLINE: usize = 8;
5const LINEAR_CONTRACTION_SCAN_LIMIT: usize = 64;
6
7/// Small axis list used by contraction preparation.
8pub(crate) type AxisVec = TinyVec<[usize; SMALL_CONTRACTION_INLINE]>;
9/// Index list used by contraction preparation.
10///
11/// `tinyvec`'s inline storage requires `I: Default`, which `IndexLike`
12/// deliberately does not require (`DynIndex` and the test index types have no
13/// cheap default). A plain `Vec` keeps the abstraction available to every
14/// `IndexLike`; the rank-bounded inline win is not reachable for a generic
15/// index type.
16pub(crate) type IndexVec<I> = Vec<I>;
17
18type AxisPairVec = TinyVec<[(usize, usize); SMALL_CONTRACTION_INLINE]>;
19type BoolVec = TinyVec<[bool; SMALL_CONTRACTION_INLINE]>;
20
21/// Error type for index replacement operations.
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub enum ReplaceIndsError {
24    /// The symmetry space of the replacement index does not match the original.
25    SpaceMismatch {
26        /// The dimension/size of the original index
27        from_dim: usize,
28        /// The dimension/size of the replacement index
29        to_dim: usize,
30    },
31    /// Duplicate indices found in the collection.
32    DuplicateIndices {
33        /// The position of the first duplicate index
34        first_pos: usize,
35        /// The position of the duplicate index
36        duplicate_pos: usize,
37    },
38}
39
40impl std::fmt::Display for ReplaceIndsError {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        match self {
43            ReplaceIndsError::SpaceMismatch { from_dim, to_dim } => {
44                write!(
45                    f,
46                    "Index space mismatch: cannot replace index with dimension {} with index of dimension {}",
47                    from_dim, to_dim
48                )
49            }
50            ReplaceIndsError::DuplicateIndices {
51                first_pos,
52                duplicate_pos,
53            } => {
54                write!(
55                    f,
56                    "Duplicate indices found: index at position {} equals index at position {}",
57                    duplicate_pos, first_pos
58                )
59            }
60        }
61    }
62}
63
64impl std::error::Error for ReplaceIndsError {}
65
66/// Check if a collection of indices contains any duplicate full indices.
67///
68/// # Arguments
69/// * `indices` - Collection of indices to check
70///
71/// # Returns
72/// `Ok(())` if all indices are unique, or `Err(ReplaceIndsError::DuplicateIndices)` if duplicates are found.
73///
74/// # Errors
75///
76/// Returns [`ReplaceIndsError::DuplicateIndices`] when two entries are equal,
77/// naming the first and duplicate positions.
78///
79/// # Example
80/// ```
81/// use tensor4all_core::index::{DefaultIndex as Index, DynId};
82/// use tensor4all_core::index_ops::check_unique_indices;
83///
84/// let i = Index::new_dyn(2);
85/// let j = Index::new_dyn(3);
86/// let indices = vec![i.clone(), j.clone()];
87/// assert!(check_unique_indices(&indices).is_ok());
88///
89/// let duplicate = vec![i.clone(), i.clone()];
90/// assert!(check_unique_indices(&duplicate).is_err());
91/// ```
92pub fn check_unique_indices<I: IndexLike>(indices: &[I]) -> Result<(), ReplaceIndsError> {
93    use std::collections::HashMap;
94    let mut seen: HashMap<&I, usize> = HashMap::with_capacity(indices.len());
95    for (pos, idx) in indices.iter().enumerate() {
96        if let Some(&first_pos) = seen.get(idx) {
97            return Err(ReplaceIndsError::DuplicateIndices {
98                first_pos,
99                duplicate_pos: pos,
100            });
101        }
102        seen.insert(idx, pos);
103    }
104    Ok(())
105}
106
107/// Replace indices in a collection based on full-index matching.
108///
109/// This corresponds to ITensors.jl's `replace_indices` function. It replaces indices
110/// in `indices` that equal any of the `(old, new)` pairs in `replacements`.
111/// The replacement index must have the same dimension as the original.
112///
113/// # Arguments
114/// * `indices` - Collection of indices to modify
115/// * `replacements` - Pairs of `(old_index, new_index)` where indices equal to `old_index` are replaced with `new_index`
116///
117/// # Returns
118/// A new vector with replacements applied, or an error if any replacement has a shape mismatch.
119///
120/// # Errors
121/// Returns `ReplaceIndsError::SpaceMismatch` if any replacement index has a different dimension than the original.
122///
123/// # Example
124/// ```
125/// use tensor4all_core::index::{DefaultIndex as Index, DynId};
126/// use tensor4all_core::index_ops::replace_indices;
127///
128/// let i = Index::new_dyn(2);
129/// let j = Index::new_dyn(3);
130/// let k = Index::new_dyn(4);
131/// let new_j = Index::new_dyn(3);  // Same size as j
132///
133/// let indices = vec![i.clone(), j.clone(), k.clone()];
134/// let replacements = vec![(j.clone(), new_j.clone())];
135///
136/// let replaced = replace_indices(indices, &replacements).unwrap();
137/// assert_eq!(replaced.len(), 3);
138/// assert_eq!(replaced[1].id, new_j.id);
139/// ```
140pub fn replace_indices<I: IndexLike>(
141    indices: Vec<I>,
142    replacements: &[(I, I)],
143) -> Result<Vec<I>, ReplaceIndsError> {
144    // Check for duplicates in input indices
145    check_unique_indices(&indices)?;
146
147    // Build a map from old index to new index for fast lookup.
148    let mut replacement_map = std::collections::HashMap::with_capacity(replacements.len());
149    for (old, new) in replacements {
150        // Validate dimension match
151        if old.dim() != new.dim() {
152            return Err(ReplaceIndsError::SpaceMismatch {
153                from_dim: old.dim(),
154                to_dim: new.dim(),
155            });
156        }
157        replacement_map.insert(old.clone(), new);
158    }
159
160    // Apply replacements
161    let mut result = Vec::with_capacity(indices.len());
162    for idx in indices {
163        if let Some(new_idx) = replacement_map.get(&idx) {
164            result.push((*new_idx).clone());
165        } else {
166            result.push(idx);
167        }
168    }
169
170    // Check for duplicates in result indices
171    check_unique_indices(&result)?;
172    Ok(result)
173}
174
175/// Replace indices in-place based on full-index matching.
176///
177/// This is an in-place variant of `replace_indices` that modifies the input slice directly.
178/// Useful for performance-critical code where you want to avoid allocations.
179///
180/// # Arguments
181/// * `indices` - Mutable slice of indices to modify
182/// * `replacements` - Pairs of `(old_index, new_index)` where indices equal to `old_index` are replaced with `new_index`
183///
184/// # Returns
185/// `Ok(())` on success, or an error if any replacement has a shape mismatch.
186///
187/// # Errors
188/// Returns `ReplaceIndsError::SpaceMismatch` if any replacement index has a different dimension than the original.
189///
190/// # Example
191/// ```
192/// use tensor4all_core::index::{DefaultIndex as Index, DynId};
193/// use tensor4all_core::index_ops::replace_indices_mut;
194///
195/// let i = Index::new_dyn(2);
196/// let j = Index::new_dyn(3);
197/// let k = Index::new_dyn(4);
198/// let new_j = Index::new_dyn(3);
199///
200/// let mut indices = vec![i.clone(), j.clone(), k.clone()];
201/// let replacements = vec![(j.clone(), new_j.clone())];
202///
203/// replace_indices_mut(&mut indices, &replacements).unwrap();
204/// assert_eq!(indices[1].id, new_j.id);
205/// ```
206pub fn replace_indices_mut<I: IndexLike>(
207    indices: &mut [I],
208    replacements: &[(I, I)],
209) -> Result<(), ReplaceIndsError> {
210    // Check for duplicates in input indices
211    check_unique_indices(indices)?;
212
213    // Build a map from old index to new index for fast lookup.
214    let mut replacement_map = std::collections::HashMap::with_capacity(replacements.len());
215    for (old, new) in replacements {
216        // Validate dimension match
217        if old.dim() != new.dim() {
218            return Err(ReplaceIndsError::SpaceMismatch {
219                from_dim: old.dim(),
220                to_dim: new.dim(),
221            });
222        }
223        replacement_map.insert(old.clone(), new);
224    }
225
226    // Stage the replacements so every validation completes before the caller's
227    // slice is modified.
228    let mut result = indices.to_vec();
229    for idx in &mut result {
230        if let Some(new_idx) = replacement_map.get(idx) {
231            *idx = (*new_idx).clone();
232        }
233    }
234
235    check_unique_indices(&result)?;
236    indices.clone_from_slice(&result);
237    Ok(())
238}
239
240/// Find indices that are unique to the first collection (set difference A \ B).
241///
242/// Returns indices that appear in `indices_a` but not in `indices_b` (matched by full index).
243/// This corresponds to ITensors.jl's `uniqueinds` function.
244///
245/// # Arguments
246/// * `indices_a` - First collection of indices
247/// * `indices_b` - Second collection of indices
248///
249/// # Returns
250/// A vector containing indices from `indices_a` that are not in `indices_b`.
251///
252/// # Example
253/// ```
254/// use tensor4all_core::index::{DefaultIndex as Index, DynId};
255/// use tensor4all_core::index_ops::unique_inds;
256///
257/// let i = Index::new_dyn(2);
258/// let j = Index::new_dyn(3);
259/// let k = Index::new_dyn(4);
260///
261/// let indices_a = vec![i.clone(), j.clone()];
262/// let indices_b = vec![j.clone(), k.clone()];
263///
264/// let unique = unique_inds(&indices_a, &indices_b);
265/// assert_eq!(unique.len(), 1);
266/// assert_eq!(unique[0].id, i.id);
267/// ```
268pub fn unique_inds<I: IndexLike>(indices_a: &[I], indices_b: &[I]) -> Vec<I> {
269    indices_a
270        .iter()
271        .filter(|idx| !indices_b.iter().any(|other| other == *idx))
272        .cloned()
273        .collect()
274}
275
276/// Find indices that are not common between two collections (symmetric difference).
277///
278/// Returns indices that appear in either `indices_a` or `indices_b` but not in both
279/// (matched by full index). This corresponds to ITensors.jl's `noncommoninds` function.
280///
281/// Time complexity: O(n + m) where n = len(indices_a), m = len(indices_b).
282///
283/// # Arguments
284/// * `indices_a` - First collection of indices
285/// * `indices_b` - Second collection of indices
286///
287/// # Returns
288/// A vector containing indices from both collections that are not common to both.
289/// Order: indices from A first (in original order), then indices from B (in original order).
290///
291/// # Example
292/// ```
293/// use tensor4all_core::index::{DefaultIndex as Index, DynId};
294/// use tensor4all_core::index_ops::noncommon_inds;
295///
296/// let i = Index::new_dyn(2);
297/// let j = Index::new_dyn(3);
298/// let k = Index::new_dyn(4);
299///
300/// let indices_a = vec![i.clone(), j.clone()];
301/// let indices_b = vec![j.clone(), k.clone()];
302///
303/// let noncommon = noncommon_inds(&indices_a, &indices_b);
304/// assert_eq!(noncommon.len(), 2);  // i and k
305/// ```
306pub fn noncommon_inds<I: IndexLike>(indices_a: &[I], indices_b: &[I]) -> Vec<I> {
307    // Pre-allocate with estimated capacity (worst case: no common indices)
308    let mut result = Vec::with_capacity(indices_a.len() + indices_b.len());
309
310    // Add indices from A that are not in B
311    result.extend(
312        indices_a
313            .iter()
314            .filter(|idx| !indices_b.iter().any(|other| other == *idx))
315            .cloned(),
316    );
317    // Add indices from B that are not in A
318    result.extend(
319        indices_b
320            .iter()
321            .filter(|idx| !indices_a.iter().any(|other| other == *idx))
322            .cloned(),
323    );
324    result
325}
326
327/// Find the union of two index collections.
328///
329/// Returns all unique indices from both collections (matched by full index).
330/// This corresponds to ITensors.jl's `unioninds` function.
331///
332/// Time complexity: O(n + m) where n = len(indices_a), m = len(indices_b).
333///
334/// # Arguments
335/// * `indices_a` - First collection of indices
336/// * `indices_b` - Second collection of indices
337///
338/// # Returns
339/// A vector containing all unique indices from both collections.
340///
341/// # Example
342/// ```
343/// use tensor4all_core::index::{DefaultIndex as Index, DynId};
344/// use tensor4all_core::index_ops::union_inds;
345///
346/// let i = Index::new_dyn(2);
347/// let j = Index::new_dyn(3);
348/// let k = Index::new_dyn(4);
349///
350/// let indices_a = vec![i.clone(), j.clone()];
351/// let indices_b = vec![j.clone(), k.clone()];
352///
353/// let union = union_inds(&indices_a, &indices_b);
354/// assert_eq!(union.len(), 3);  // i, j, k
355/// ```
356pub fn union_inds<I: IndexLike>(indices_a: &[I], indices_b: &[I]) -> Vec<I> {
357    let mut seen: std::collections::HashSet<&I> =
358        std::collections::HashSet::with_capacity(indices_a.len() + indices_b.len());
359    let mut result = Vec::with_capacity(indices_a.len() + indices_b.len());
360
361    for idx in indices_a {
362        if seen.insert(idx) {
363            result.push(idx.clone());
364        }
365    }
366    for idx in indices_b {
367        if seen.insert(idx) {
368            result.push(idx.clone());
369        }
370    }
371    result
372}
373
374/// Check if a collection contains a specific full index.
375///
376/// This corresponds to ITensors.jl's `hasind` function.
377///
378/// # Arguments
379/// * `indices` - Collection of indices to search
380/// * `index` - The index to look for
381///
382/// # Returns
383/// `true` if an index with matching ID is found, `false` otherwise.
384///
385/// # Example
386/// ```
387/// use tensor4all_core::index::{DefaultIndex as Index, DynId};
388/// use tensor4all_core::index_ops::hasind;
389///
390/// let i = Index::new_dyn(2);
391/// let j = Index::new_dyn(3);
392/// let indices = vec![i.clone(), j.clone()];
393///
394/// assert!(hasind(&indices, &i));
395/// assert!(!hasind(&indices, &Index::new_dyn(4)));
396/// ```
397pub fn hasind<I: IndexLike>(indices: &[I], index: &I) -> bool {
398    indices.iter().any(|idx| idx == index)
399}
400
401/// Check if a collection contains all of the specified full indices.
402///
403/// This corresponds to ITensors.jl's `hasinds` function.
404///
405/// # Arguments
406/// * `indices` - Collection of indices to search
407/// * `targets` - The indices to look for
408///
409/// # Returns
410/// `true` if all target indices are found, `false` otherwise.
411///
412/// # Example
413/// ```
414/// use tensor4all_core::index::{DefaultIndex as Index, DynId};
415/// use tensor4all_core::index_ops::has_inds;
416///
417/// let i = Index::new_dyn(2);
418/// let j = Index::new_dyn(3);
419/// let k = Index::new_dyn(4);
420/// let indices = vec![i.clone(), j.clone(), k.clone()];
421///
422/// assert!(has_inds(&indices, &[i.clone(), j.clone()]));
423/// assert!(!has_inds(&indices, &[i.clone(), Index::new_dyn(5)]));
424/// ```
425#[doc(alias = "hasinds")]
426pub fn has_inds<I: IndexLike>(indices: &[I], targets: &[I]) -> bool {
427    targets
428        .iter()
429        .all(|target| indices.iter().any(|idx| idx == target))
430}
431
432/// Check if two collections have any common full indices.
433///
434/// This corresponds to ITensors.jl's `hascommoninds` function.
435///
436/// # Arguments
437/// * `indices_a` - First collection of indices
438/// * `indices_b` - Second collection of indices
439///
440/// # Returns
441/// `true` if there is at least one common index, `false` otherwise.
442///
443/// # Example
444/// ```
445/// use tensor4all_core::index::{DefaultIndex as Index, DynId};
446/// use tensor4all_core::index_ops::has_common_inds;
447///
448/// let i = Index::new_dyn(2);
449/// let j = Index::new_dyn(3);
450/// let k = Index::new_dyn(4);
451///
452/// let indices_a = vec![i.clone(), j.clone()];
453/// let indices_b = vec![j.clone(), k.clone()];
454///
455/// assert!(has_common_inds(&indices_a, &indices_b));
456/// assert!(!has_common_inds(&[i.clone()], &[k.clone()]));
457/// ```
458#[doc(alias = "hascommoninds")]
459pub fn has_common_inds<I: IndexLike>(indices_a: &[I], indices_b: &[I]) -> bool {
460    indices_a
461        .iter()
462        .any(|idx| indices_b.iter().any(|other| other == idx))
463}
464
465/// Find common indices between two index collections.
466///
467/// Returns a vector of indices that appear in both `indices_a` and `indices_b`
468/// (set intersection). This is similar to ITensors.jl's `commoninds` function.
469///
470/// Time complexity: O(n + m) where n = len(indices_a), m = len(indices_b).
471///
472/// # Arguments
473/// * `indices_a` - First collection of indices
474/// * `indices_b` - Second collection of indices
475///
476/// # Returns
477/// A vector containing indices that are common to both collections (matched by full index).
478///
479/// # Example
480/// ```
481/// use tensor4all_core::index::{DefaultIndex as Index, DynId};
482/// use tensor4all_core::index_ops::common_inds;
483///
484/// let i = Index::new_dyn(2);
485/// let j = Index::new_dyn(3);
486/// let k = Index::new_dyn(4);
487///
488/// let indices_a = vec![i.clone(), j.clone()];
489/// let indices_b = vec![j.clone(), k.clone()];
490///
491/// let common = common_inds(&indices_a, &indices_b);
492/// assert_eq!(common.len(), 1);
493/// assert_eq!(common[0].id, j.id);
494/// ```
495pub fn common_inds<I: IndexLike>(indices_a: &[I], indices_b: &[I]) -> Vec<I> {
496    indices_a
497        .iter()
498        .filter(|idx| indices_b.iter().any(|other| other == *idx))
499        .cloned()
500        .collect()
501}
502
503/// Find contractable indices between two slices and return their positions.
504///
505/// Returns a vector of `(pos_a, pos_b)` tuples where each tuple indicates
506/// that `indices_a[pos_a]` and `indices_b[pos_b]` are contractable
507/// (same ID, same dimension, and compatible ConjState).
508///
509/// # Example
510/// ```
511/// use tensor4all_core::index::DefaultIndex as Index;
512/// use tensor4all_core::index_ops::common_ind_positions;
513///
514/// let i = Index::new_dyn(2);
515/// let j = Index::new_dyn(3);
516/// let k = Index::new_dyn(4);
517///
518/// let indices_a = vec![i.clone(), j.clone()];
519/// let indices_b = vec![j.clone(), k.clone()];
520///
521/// let positions = common_ind_positions(&indices_a, &indices_b);
522/// assert_eq!(positions, vec![(1, 0)]); // j is at position 1 in a, position 0 in b
523/// ```
524pub fn common_ind_positions<I: IndexLike>(indices_a: &[I], indices_b: &[I]) -> Vec<(usize, usize)> {
525    common_ind_positions_small(indices_a, indices_b).into_vec()
526}
527
528fn common_ind_positions_small<I: IndexLike>(indices_a: &[I], indices_b: &[I]) -> AxisPairVec {
529    let scan_work = indices_a.len().saturating_mul(indices_b.len());
530    if scan_work <= LINEAR_CONTRACTION_SCAN_LIMIT {
531        return common_ind_positions_linear(indices_a, indices_b);
532    }
533    common_ind_positions_hashed(indices_a, indices_b)
534}
535
536fn common_ind_positions_linear<I: IndexLike>(indices_a: &[I], indices_b: &[I]) -> AxisPairVec {
537    let mut positions = AxisPairVec::new();
538    for (pos_a, idx_a) in indices_a.iter().enumerate() {
539        for (pos_b, idx_b) in indices_b.iter().enumerate() {
540            if idx_a.is_contractable(idx_b) {
541                positions.push((pos_a, pos_b));
542                break; // Each index in a can match at most one in b
543            }
544        }
545    }
546    positions
547}
548
549fn common_ind_positions_hashed<I: IndexLike>(indices_a: &[I], indices_b: &[I]) -> AxisPairVec {
550    use std::collections::HashMap;
551
552    // Bucket indices_b by their FULL value and probe each `idx_a` under the
553    // key `conj(idx_a)`: the linear scan accepts exactly `b == conj(a)`
554    // (undirected: conj is the identity; directed: is_contractable requires
555    // `conj(a) == b`, even when conj changes the id). Keying by Id, or by
556    // bare value with a same-value probe, would silently drop directed pairs
557    // whose conj partner carries a different id (#615).
558    let mut positions_by_value: HashMap<I, TinyVec<[usize; 2]>> =
559        HashMap::with_capacity(indices_b.len());
560    for (pos_b, idx_b) in indices_b.iter().enumerate() {
561        positions_by_value
562            .entry(idx_b.clone())
563            .or_default()
564            .push(pos_b);
565    }
566
567    let mut positions = AxisPairVec::new();
568    for (pos_a, idx_a) in indices_a.iter().enumerate() {
569        let Some(candidate_positions) = positions_by_value.get(&idx_a.conj()) else {
570            continue;
571        };
572        for &pos_b in candidate_positions {
573            if idx_a.is_contractable(&indices_b[pos_b]) {
574                positions.push((pos_a, pos_b));
575                break; // Each index in a can match at most one in b
576            }
577        }
578    }
579    positions
580}
581
582/// Result of preparing a tensor contraction.
583///
584/// Contains all the information needed to perform the contraction:
585/// - Which axes to contract from each tensor
586/// - The resulting indices after contraction
587#[derive(Debug, Clone)]
588pub(crate) struct ContractionSpec<I: IndexLike> {
589    /// Axes to contract from the first tensor (positions in `indices_a`).
590    pub axes_a: AxisVec,
591    /// Axes to contract from the second tensor (positions in `indices_b`).
592    pub axes_b: AxisVec,
593    /// Indices of the result tensor (non-contracted indices from both tensors).
594    pub result_indices: IndexVec<I>,
595}
596
597/// Error type for contraction preparation.
598#[derive(Debug, Clone, PartialEq, Eq)]
599pub(crate) enum ContractionError {
600    /// No common indices found for contraction.
601    NoCommonIndices,
602    /// Dimension mismatch for a common index.
603    DimensionMismatch {
604        /// Position in the first tensor.
605        pos_a: usize,
606        /// Position in the second tensor.
607        pos_b: usize,
608        /// Dimension in the first tensor.
609        dim_a: usize,
610        /// Dimension in the second tensor.
611        dim_b: usize,
612    },
613    /// Duplicate axis specified in contraction.
614    DuplicateAxis {
615        /// Which tensor has the duplicate ("self" or "other").
616        tensor: &'static str,
617        /// Position of the duplicate axis.
618        pos: usize,
619    },
620    /// Index not found in tensor.
621    IndexNotFound {
622        /// Which tensor the index was not found in.
623        tensor: &'static str,
624    },
625    /// Batch contraction not yet implemented.
626    BatchContractionNotImplemented,
627}
628
629impl std::fmt::Display for ContractionError {
630    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
631        match self {
632            ContractionError::NoCommonIndices => {
633                write!(f, "No common indices found for contraction")
634            }
635            ContractionError::DimensionMismatch {
636                pos_a,
637                pos_b,
638                dim_a,
639                dim_b,
640            } => {
641                write!(
642                    f,
643                    "Dimension mismatch: tensor_a[{}]={} != tensor_b[{}]={}",
644                    pos_a, dim_a, pos_b, dim_b
645                )
646            }
647            ContractionError::DuplicateAxis { tensor, pos } => {
648                write!(f, "Duplicate axis {} in {} tensor", pos, tensor)
649            }
650            ContractionError::IndexNotFound { tensor } => {
651                write!(f, "Index not found in {} tensor", tensor)
652            }
653            ContractionError::BatchContractionNotImplemented => {
654                write!(f, "Batch contraction not yet implemented")
655            }
656        }
657    }
658}
659
660impl std::error::Error for ContractionError {}
661
662/// Prepare contraction data for two tensors that share common indices.
663///
664/// This internal helper finds common indices and computes the axes to contract
665/// together with the resulting non-contracted indices.
666pub(crate) fn prepare_contraction<I: IndexLike>(
667    indices_a: &[I],
668    dims_a: &[usize],
669    indices_b: &[I],
670    dims_b: &[usize],
671) -> Result<ContractionSpec<I>, ContractionError> {
672    // Find common indices and their positions.
673    // If no common indices exist, this becomes an outer product (empty axes).
674    let positions = common_ind_positions_small(indices_a, indices_b);
675
676    let mut axes_a = AxisVec::with_capacity(positions.len());
677    let mut axes_b = AxisVec::with_capacity(positions.len());
678    for &(pos_a, pos_b) in &positions {
679        axes_a.push(pos_a);
680        axes_b.push(pos_b);
681    }
682
683    // Verify dimensions match
684    for &(pos_a, pos_b) in &positions {
685        if dims_a[pos_a] != dims_b[pos_b] {
686            return Err(ContractionError::DimensionMismatch {
687                pos_a,
688                pos_b,
689                dim_a: dims_a[pos_a],
690                dim_b: dims_b[pos_b],
691            });
692        }
693    }
694
695    let result_indices = build_contraction_result_indices(indices_a, &axes_a, indices_b, &axes_b);
696
697    Ok(ContractionSpec {
698        axes_a,
699        axes_b,
700        result_indices,
701    })
702}
703
704/// Prepare contraction data for explicit index pairs (like tensordot).
705///
706/// Unlike `prepare_contraction`, this internal helper takes explicit pairs of
707/// indices to contract, allowing contraction of indices with different IDs.
708pub(crate) fn prepare_contraction_pairs<I: IndexLike>(
709    indices_a: &[I],
710    dims_a: &[usize],
711    indices_b: &[I],
712    dims_b: &[usize],
713    pairs: &[(I, I)],
714) -> Result<ContractionSpec<I>, ContractionError> {
715    if pairs.is_empty() {
716        return Err(ContractionError::NoCommonIndices);
717    }
718
719    // Check for batch contraction (common indices not in pairs). The explicit
720    // pair list identifies axes by full index metadata, not by ID alone.
721    let common_positions = common_ind_positions_small(indices_a, indices_b);
722    for (pos_a, pos_b) in &common_positions {
723        let idx_a = &indices_a[*pos_a];
724        let idx_b = &indices_b[*pos_b];
725        if !pairs
726            .iter()
727            .any(|(contracted_idx, _)| contracted_idx == idx_a)
728            || !pairs
729                .iter()
730                .any(|(_, contracted_idx)| contracted_idx == idx_b)
731        {
732            return Err(ContractionError::BatchContractionNotImplemented);
733        }
734    }
735
736    // Find positions and validate
737    let mut axes_a = AxisVec::with_capacity(pairs.len());
738    let mut axes_b = AxisVec::with_capacity(pairs.len());
739    let mut contracted_a = bool_flags(indices_a.len());
740    let mut contracted_b = bool_flags(indices_b.len());
741
742    for (idx_a, idx_b) in pairs {
743        let pos_a = indices_a
744            .iter()
745            .position(|idx| idx == idx_a)
746            .ok_or(ContractionError::IndexNotFound { tensor: "self" })?;
747
748        let pos_b = indices_b
749            .iter()
750            .position(|idx| idx == idx_b)
751            .ok_or(ContractionError::IndexNotFound { tensor: "other" })?;
752
753        // Verify dimensions match
754        if dims_a[pos_a] != dims_b[pos_b] {
755            return Err(ContractionError::DimensionMismatch {
756                pos_a,
757                pos_b,
758                dim_a: dims_a[pos_a],
759                dim_b: dims_b[pos_b],
760            });
761        }
762
763        // Check for duplicate axes
764        if contracted_a[pos_a] {
765            return Err(ContractionError::DuplicateAxis {
766                tensor: "self",
767                pos: pos_a,
768            });
769        }
770        if contracted_b[pos_b] {
771            return Err(ContractionError::DuplicateAxis {
772                tensor: "other",
773                pos: pos_b,
774            });
775        }
776
777        contracted_a[pos_a] = true;
778        contracted_b[pos_b] = true;
779        axes_a.push(pos_a);
780        axes_b.push(pos_b);
781    }
782
783    let result_indices = build_contraction_result_indices(indices_a, &axes_a, indices_b, &axes_b);
784
785    Ok(ContractionSpec {
786        axes_a,
787        axes_b,
788        result_indices,
789    })
790}
791
792fn bool_flags(len: usize) -> BoolVec {
793    let mut flags = BoolVec::with_capacity(len);
794    flags.resize(len, false);
795    flags
796}
797
798fn build_contraction_result_indices<I: IndexLike>(
799    indices_a: &[I],
800    axes_a: &[usize],
801    indices_b: &[I],
802    axes_b: &[usize],
803) -> IndexVec<I> {
804    let mut contracted_a = bool_flags(indices_a.len());
805    let mut contracted_b = bool_flags(indices_b.len());
806    for &axis in axes_a {
807        contracted_a[axis] = true;
808    }
809    for &axis in axes_b {
810        contracted_b[axis] = true;
811    }
812
813    let result_len = indices_a.len() + indices_b.len() - axes_a.len() - axes_b.len();
814    let mut result_indices = IndexVec::with_capacity(result_len);
815
816    for (i, idx) in indices_a.iter().enumerate() {
817        if !contracted_a[i] {
818            result_indices.push(idx.clone());
819        }
820    }
821
822    for (i, idx) in indices_b.iter().enumerate() {
823        if !contracted_b[i] {
824            result_indices.push(idx.clone());
825        }
826    }
827
828    result_indices
829}
830
831#[cfg(test)]
832mod tests;
tensor4all_core/index_ops.rs

tensor4all_core/
index_ops.rs
