tensor4all_core/
index_ops.rs

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

tensor4all_core/
index_ops.rs
