Skip to main content

tensor4all_core/defaults/
index.rs

1//! Index types for tensor network operations.
2//!
3//! This module provides the default index types:
4//!
5//! - [`DynId`]: Runtime identity (UUID-based unique identifier)
6//! - [`TagSet`]: Tag set for metadata (Arc-wrapped for cheap cloning)
7//! - [`Index`]: Generic index type parameterized by Id and Tags
8//! - [`DynIndex`]: Default index type (`Index<DynId, TagSet>`)
9//!
10//! The `DynIndex` type implements the [`IndexLike`] trait.
11//!
12//! **Note**: Symmetry (quantum numbers) is not included in the default implementation.
13//! For QSpace-compatible indices with non-Abelian symmetries, use a separate concrete type
14//! that implements `IndexLike` directly.
15
16use crate::index_like::IndexLike;
17use crate::tagset::{DefaultTagSet as InlineTagSet, TagSetError, TagSetIterator, TagSetLike};
18use anyhow::Result;
19use rand::Rng;
20use std::cell::RefCell;
21use std::sync::Arc;
22
23/// Runtime ID for ITensors-like dynamic identity.
24///
25/// Uses UInt64 for compatibility with ITensors.jl's `IDType = UInt64`.
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
27pub struct DynId(pub u64);
28
29/// Tag set wrapper using `Arc` for efficient cloning.
30///
31/// This wraps the underlying tag storage in an `Arc` for cheap cloning (reference count increment only).
32/// Tags are immutable and shared across indices with the same tag set.
33///
34/// # Size comparison
35/// - Inline storage: 168 bytes (Copy)
36/// - `TagSet` (Arc): 8 bytes (Clone only)
37///
38/// # Example
39/// ```
40/// use tensor4all_core::index::TagSet;
41///
42/// let tags = TagSet::from_str("Site,Link").unwrap();
43/// assert!(tags.has_tag("Site"));
44/// assert!(tags.has_tag("Link"));
45/// ```
46#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
47pub struct TagSet(Arc<InlineTagSet>);
48
49impl TagSet {
50    /// Create an empty tag set.
51    pub fn new() -> Self {
52        Self(Arc::new(InlineTagSet::new()))
53    }
54
55    /// Create a tag set from a comma-separated string.
56    #[allow(clippy::should_implement_trait)]
57    pub fn from_str(s: &str) -> Result<Self, TagSetError> {
58        Ok(Self(Arc::new(InlineTagSet::from_str(s)?)))
59    }
60
61    /// Create a tag set from a slice of tag strings.
62    ///
63    /// Returns an error if any tag contains a comma (reserved as separator in `from_str`).
64    ///
65    /// # Example
66    /// ```
67    /// use tensor4all_core::index::TagSet;
68    ///
69    /// let tags = TagSet::from_tags(&["Site", "Link"]).unwrap();
70    /// assert!(tags.has_tag("Site"));
71    /// assert!(tags.has_tag("Link"));
72    /// assert_eq!(tags.len(), 2);
73    ///
74    /// // Comma in tag is an error
75    /// assert!(TagSet::from_tags(&["Site,Link"]).is_err());
76    /// ```
77    pub fn from_tags(tags: &[&str]) -> Result<Self, TagSetError> {
78        let mut inner = InlineTagSet::new();
79        for tag in tags {
80            if tag.contains(',') {
81                return Err(TagSetError::TagContainsComma {
82                    tag: (*tag).to_string(),
83                });
84            }
85            inner.add_tag(tag)?;
86        }
87        Ok(Self(Arc::new(inner)))
88    }
89
90    /// Check if a tag is present.
91    pub fn has_tag(&self, tag: &str) -> bool {
92        self.0.has_tag(tag)
93    }
94
95    /// Get the number of tags.
96    pub fn len(&self) -> usize {
97        self.0.len()
98    }
99
100    /// Check if the tag set is empty.
101    pub fn is_empty(&self) -> bool {
102        self.0.is_empty()
103    }
104
105    /// Get the inner Arc for advanced use.
106    pub fn inner(&self) -> &Arc<InlineTagSet> {
107        &self.0
108    }
109}
110
111impl std::ops::Deref for TagSet {
112    type Target = InlineTagSet;
113
114    fn deref(&self) -> &Self::Target {
115        &self.0
116    }
117}
118
119impl TagSetLike for TagSet {
120    fn len(&self) -> usize {
121        self.0.len()
122    }
123
124    fn capacity(&self) -> usize {
125        self.0.capacity()
126    }
127
128    fn get(&self, index: usize) -> Option<String> {
129        TagSetLike::get(&*self.0, index)
130    }
131
132    fn iter(&self) -> TagSetIterator<'_> {
133        TagSetLike::iter(&*self.0)
134    }
135
136    fn has_tag(&self, tag: &str) -> bool {
137        self.0.has_tag(tag)
138    }
139
140    fn add_tag(&mut self, tag: &str) -> Result<(), TagSetError> {
141        // Arc is immutable, so we need to clone and replace
142        let mut inner = *self.0;
143        inner.add_tag(tag)?;
144        self.0 = Arc::new(inner);
145        Ok(())
146    }
147
148    fn remove_tag(&mut self, tag: &str) -> bool {
149        // Arc is immutable, so we need to clone and replace
150        let mut inner = *self.0;
151        let removed = inner.remove_tag(tag);
152        if removed {
153            self.0 = Arc::new(inner);
154        }
155        removed
156    }
157}
158
159/// Index with generic identity type `Id` and tag type `Tags`.
160///
161/// - `Id = DynId` for ITensors-like runtime identity
162/// - `Id = ZST marker type` for compile-time-known identity
163/// - `Tags = TagSet` for tags (default, Arc-wrapped for cheap cloning)
164///
165/// **Note**: This default implementation does not include symmetry (quantum numbers).
166/// For QSpace-compatible indices with non-Abelian symmetries, use a separate concrete type.
167///
168/// # Memory Layout
169/// With default types (`DynId`, `TagSet`):
170/// - Size: 32 bytes (8 + 8 + 8 + 8)
171/// - Tags are shared via `Arc`, so cloning is cheap (reference count increment only)
172///
173/// **Equality**: Two `Index` values are considered equal if and only if their `id` and `tags`
174/// fields and `plev` match (matching ITensors.jl semantics where equality = id + plev + tags).
175///
176/// # Example
177/// ```
178/// use tensor4all_core::index::{Index, DynId, TagSet};
179///
180/// // Create shared tags once
181/// let site_tags = TagSet::from_str("Site").unwrap();
182///
183/// // Share the same tags across many indices (cheap clone)
184/// let i1 = Index::<DynId>::new_dyn_with_tags(2, site_tags.clone());
185/// let i2 = Index::<DynId>::new_dyn_with_tags(2, site_tags.clone());
186/// // i1.tags and i2.tags point to the same Arc
187/// ```
188#[derive(Debug, Clone)]
189pub struct Index<Id, Tags = TagSet> {
190    /// The unique identifier for this index.
191    pub id: Id,
192    /// The dimension (size) of this index.
193    pub dim: usize,
194    /// The prime level of this index.
195    pub plev: i64,
196    /// The tag set associated with this index.
197    pub tags: Tags,
198}
199
200impl<Id, Tags> Index<Id, Tags>
201where
202    Tags: Default,
203{
204    /// Create a new index with the given identity and dimension.
205    pub fn new(id: Id, dim: usize) -> Self {
206        Self {
207            id,
208            dim,
209            plev: 0,
210            tags: Tags::default(),
211        }
212    }
213
214    /// Create a new index with the given identity, dimension, and tags.
215    pub fn new_with_tags(id: Id, dim: usize, tags: Tags) -> Self {
216        Self {
217            id,
218            dim,
219            plev: 0,
220            tags,
221        }
222    }
223
224    /// Get the dimension (size) of the index.
225    pub fn size(&self) -> usize {
226        self.dim
227    }
228
229    /// Get a reference to the tags.
230    pub fn tags(&self) -> &Tags {
231        &self.tags
232    }
233}
234
235impl<Id, Tags> Index<Id, Tags>
236where
237    Tags: Default,
238{
239    /// Create a new index from dimension (convenience constructor).
240    pub fn new_with_size(id: Id, size: usize) -> Self {
241        Self {
242            id,
243            dim: size,
244            plev: 0,
245            tags: Tags::default(),
246        }
247    }
248
249    /// Create a new index from dimension and tags.
250    pub fn new_with_size_and_tags(id: Id, size: usize, tags: Tags) -> Self {
251        Self {
252            id,
253            dim: size,
254            plev: 0,
255            tags,
256        }
257    }
258}
259
260// Constructors for Index with TagSet (default)
261impl Index<DynId, TagSet> {
262    /// Create a new index with a generated dynamic ID and no tags.
263    ///
264    /// This is the most common way to create indices. Each call generates
265    /// a unique random ID, so two calls to `new_dyn` with the same dimension
266    /// produce non-equal indices.
267    ///
268    /// # Examples
269    ///
270    /// ```
271    /// use tensor4all_core::{DynIndex, IndexLike};
272    ///
273    /// let i = DynIndex::new_dyn(4);
274    /// let j = DynIndex::new_dyn(4);
275    ///
276    /// assert_eq!(i.dim(), 4);
277    /// assert_ne!(i, j);  // different IDs
278    /// assert!(i.is_contractable(&i));  // same index is contractable with itself
279    /// assert!(!i.is_contractable(&j)); // different IDs
280    /// ```
281    pub fn new_dyn(size: usize) -> Self {
282        Self {
283            id: DynId(generate_id()),
284            dim: size,
285            plev: 0,
286            tags: TagSet::new(),
287        }
288    }
289
290    /// Create a new index with a generated dynamic ID and shared tags.
291    ///
292    /// This is the most efficient way to create many indices with the same tags.
293    /// The `Arc` is cloned (reference count increment only), not the underlying data.
294    ///
295    /// # Example
296    /// ```
297    /// use tensor4all_core::index::{Index, DynId, TagSet};
298    ///
299    /// let site_tags = TagSet::from_str("Site").unwrap();
300    /// let i1 = Index::<DynId>::new_dyn_with_tags(2, site_tags.clone());
301    /// let i2 = Index::<DynId>::new_dyn_with_tags(2, site_tags.clone());
302    /// ```
303    pub fn new_dyn_with_tags(size: usize, tags: TagSet) -> Self {
304        Self {
305            id: DynId(generate_id()),
306            dim: size,
307            plev: 0,
308            tags,
309        }
310    }
311
312    /// Create a new index with a generated dynamic ID and a single tag.
313    ///
314    /// This creates a new `TagSet` with the given tag.
315    /// For sharing the same tag across many indices, create the `TagSet`
316    /// once and use `new_dyn_with_tags` instead.
317    ///
318    /// # Examples
319    ///
320    /// ```
321    /// use tensor4all_core::DynIndex;
322    /// use tensor4all_core::TagSetLike;
323    ///
324    /// let site = DynIndex::new_dyn_with_tag(2, "Site").unwrap();
325    /// assert!(site.tags().has_tag("Site"));
326    /// ```
327    pub fn new_dyn_with_tag(size: usize, tag: &str) -> Result<Self, TagSetError> {
328        Ok(Self {
329            id: DynId(generate_id()),
330            dim: size,
331            plev: 0,
332            tags: TagSet::from_str(tag)?,
333        })
334    }
335
336    /// Create a new bond index with "Link" tag (for SVD, QR, etc.).
337    ///
338    /// This is a convenience method for creating bond indices commonly used in tensor
339    /// decompositions like SVD and QR factorization.
340    ///
341    /// # Examples
342    ///
343    /// ```
344    /// use tensor4all_core::DynIndex;
345    /// use tensor4all_core::TagSetLike;
346    ///
347    /// let link = DynIndex::new_link(4).unwrap();
348    /// assert!(link.tags().has_tag("Link"));
349    /// ```
350    pub fn new_link(size: usize) -> Result<Self, TagSetError> {
351        Self::new_dyn_with_tag(size, "Link")
352    }
353}
354
355// Equality and Hash implementations: compare by `id`, `tags`, and `plev`
356// (matching ITensors.jl semantics where equality = id + plev + tags)
357impl<Id: PartialEq, Tags: PartialEq> PartialEq for Index<Id, Tags> {
358    fn eq(&self, other: &Self) -> bool {
359        self.id == other.id && self.tags == other.tags && self.plev == other.plev
360    }
361}
362
363impl<Id: Eq, Tags: Eq> Eq for Index<Id, Tags> {}
364
365impl<Id: std::hash::Hash, Tags: std::hash::Hash> std::hash::Hash for Index<Id, Tags> {
366    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
367        self.id.hash(state);
368        self.tags.hash(state);
369        self.plev.hash(state);
370    }
371}
372
373// Copy implementation: Index is Copy when Id and Tags are both Copy
374impl<Id: Copy, Tags: Copy> Copy for Index<Id, Tags> {}
375
376thread_local! {
377    /// Thread-local random number generator for ID generation.
378    ///
379    /// Each thread has its own RNG, similar to ITensors.jl's task-local RNG.
380    /// This provides thread-safe ID generation without global synchronization.
381    static ID_RNG: RefCell<rand::rngs::ThreadRng> = RefCell::new(rand::rng());
382}
383
384/// Generate a unique random ID for dynamic indices (thread-safe).
385///
386/// Uses thread-local random number generator to generate UInt64 IDs,
387/// compatible with ITensors.jl's `IDType = UInt64`.
388pub(crate) fn generate_id() -> u64 {
389    ID_RNG.with(|rng| rng.borrow_mut().random())
390}
391
392/// Default Index type alias (same as `Index<Id>` with default tags).
393///
394/// This is provided for convenience and compatibility.
395pub type DefaultIndex<Id> = Index<Id, TagSet>;
396
397/// Type alias for backwards compatibility.
398pub type DefaultTagSet = TagSet;
399
400// ============================================================================
401// DynIndex: Default index type with IndexLike implementation
402// ============================================================================
403
404/// Type alias for the default index type with IndexLike bound.
405///
406/// `DynIndex` uses:
407/// - `DynId`: Dynamic identity (UUID-based unique identifier)
408/// - `TagSet`: Default tag set for metadata
409///
410/// This is the recommended index type for most tensor network applications.
411/// It does not include symmetry (quantum numbers); for QSpace-compatible indices,
412/// use a separate concrete type that implements `IndexLike` directly.
413///
414/// # Examples
415///
416/// ```
417/// use tensor4all_core::DynIndex;
418/// use tensor4all_core::index_like::IndexLike;
419///
420/// // Create a dynamic index with dimension 4
421/// let idx = DynIndex::new_dyn(4);
422/// assert_eq!(idx.dim(), 4);
423/// assert_eq!(idx.plev(), 0);
424///
425/// // Prime level manipulation
426/// let primed = idx.prime();
427/// assert_eq!(primed.plev(), 1);
428///
429/// let noprime = primed.noprime();
430/// assert_eq!(noprime.plev(), 0);
431///
432/// // Bond index creation (for SVD/QR)
433/// let bond = DynIndex::new_bond(8).unwrap();
434/// assert_eq!(bond.dim(), 8);
435/// ```
436pub type DynIndex = Index<DynId, TagSet>;
437
438impl IndexLike for DynIndex {
439    type Id = DynId;
440
441    fn id(&self) -> &Self::Id {
442        &self.id
443    }
444
445    fn dim(&self) -> usize {
446        self.dim
447    }
448
449    fn plev(&self) -> i64 {
450        self.plev
451    }
452
453    fn conj_state(&self) -> crate::ConjState {
454        // Default indices are undirected (ITensors.jl-like behavior)
455        crate::ConjState::Undirected
456    }
457
458    fn conj(&self) -> Self {
459        // For undirected indices, conj() is a no-op
460        self.clone()
461    }
462
463    fn sim(&self) -> Self {
464        Index {
465            id: DynId(generate_id()),
466            dim: self.dim,
467            plev: self.plev,
468            tags: self.tags.clone(),
469        }
470    }
471
472    fn create_dummy_link_pair() -> (Self, Self) {
473        let id = DynId(generate_id());
474        let idx1 = Index {
475            id,
476            dim: 1,
477            plev: 0,
478            tags: TagSet::default(),
479        };
480        let idx2 = Index {
481            id,
482            dim: 1,
483            plev: 0,
484            tags: TagSet::default(),
485        };
486        (idx1, idx2)
487    }
488
489    fn product_link(indices: &[Self]) -> Result<Self> {
490        anyhow::ensure!(
491            !indices.is_empty(),
492            "product_link requires at least one index"
493        );
494        let dim = indices.iter().try_fold(1usize, |acc, idx| {
495            acc.checked_mul(idx.dim())
496                .ok_or_else(|| anyhow::anyhow!("product link dimension overflow"))
497        })?;
498        DynIndex::new_bond(dim)
499    }
500}
501
502impl DynIndex {
503    /// Create a new bond index with a fresh identity and the specified dimension.
504    ///
505    /// This is used by factorization operations (SVD, QR) to create new internal
506    /// bond indices connecting the factors.
507    ///
508    /// # Arguments
509    /// * `dim` - The dimension of the new index
510    ///
511    /// # Returns
512    /// A new index with a unique identity and the specified dimension.
513    ///
514    /// # Examples
515    ///
516    /// ```
517    /// use tensor4all_core::{DynIndex, IndexLike};
518    ///
519    /// let bond = DynIndex::new_bond(8).unwrap();
520    /// assert_eq!(bond.dim(), 8);
521    /// ```
522    pub fn new_bond(dim: usize) -> Result<Self> {
523        Index::new_link(dim).map_err(|e| anyhow::anyhow!("Failed to create bond index: {:?}", e))
524    }
525
526    /// Return a copy of this index with its prime level incremented by one.
527    ///
528    /// Primed indices are commonly used to distinguish bra and ket indices
529    /// in MPO operations (e.g., `i` for input, `i'` for output).
530    ///
531    /// # Examples
532    ///
533    /// ```
534    /// use tensor4all_core::{DynIndex, IndexLike};
535    ///
536    /// let i = DynIndex::new_dyn(4);
537    /// assert_eq!(i.plev(), 0);
538    ///
539    /// let i_prime = i.prime();
540    /// assert_eq!(i_prime.plev(), 1);
541    ///
542    /// // Primed and unprimed indices have the same ID but are not equal
543    /// assert!(i.same_id(&i_prime));
544    /// assert_ne!(i, i_prime);
545    ///
546    /// // Primed indices are not contractable with unprimed ones
547    /// assert!(!i.is_contractable(&i_prime));
548    /// ```
549    pub fn prime(&self) -> Self {
550        let mut idx = self.clone();
551        idx.plev += 1;
552        idx
553    }
554
555    /// Return a copy of this index with its prime level reset to zero.
556    ///
557    /// # Examples
558    ///
559    /// ```
560    /// use tensor4all_core::{DynIndex, IndexLike};
561    ///
562    /// let i = DynIndex::new_dyn(4);
563    /// let i2 = i.prime().prime();
564    /// assert_eq!(i2.plev(), 2);
565    ///
566    /// let i0 = i2.noprime();
567    /// assert_eq!(i0.plev(), 0);
568    /// assert_eq!(i0, i);
569    /// ```
570    pub fn noprime(&self) -> Self {
571        let mut idx = self.clone();
572        idx.plev = 0;
573        idx
574    }
575
576    /// Return a copy of this index with its prime level set explicitly.
577    ///
578    /// # Examples
579    ///
580    /// ```
581    /// use tensor4all_core::{DynIndex, IndexLike};
582    ///
583    /// let i = DynIndex::new_dyn(4);
584    /// let i3 = i.set_plev(3);
585    /// assert_eq!(i3.plev(), 3);
586    /// ```
587    pub fn set_plev(&self, plev: i64) -> Self {
588        let mut idx = self.clone();
589        idx.plev = plev;
590        idx
591    }
592}
593
594#[cfg(test)]
595mod tests;