tensor4all_core/
block_tensor.rs

1//! Block tensor type for GMRES with block matrices.
2//!
3//! This module provides [`BlockTensor`], a collection of tensors organized
4//! in a block structure. It implements [`TensorLike`] for the vector space
5//! operations required by GMRES, allowing block matrix linear equations
6//! `Ax = b` to be solved using the existing GMRES implementation.
7//!
8//! # Example
9//!
10//! ```
11//! use tensor4all_core::block_tensor::BlockTensor;
12//! use tensor4all_core::krylov::{gmres, GmresOptions};
13//! use tensor4all_core::{DynIndex, IdxTensor};
14//!
15//! # fn main() -> anyhow::Result<()> {
16//! let i = DynIndex::new_dyn(2);
17//! let b_block = IdxTensor::from_dense(vec![i.clone()], vec![1.0, 2.0])?;
18//! let zero_block = IdxTensor::from_dense(vec![i.clone()], vec![0.0, 0.0])?;
19//!
20//! let b = BlockTensor::new(vec![b_block], (1, 1))?;
21//! let x0 = BlockTensor::new(vec![zero_block], (1, 1))?;
22//!
23//! let apply_a = |x: &BlockTensor<IdxTensor>| Ok(x.clone());
24//! let result = gmres(apply_a, &b, &x0, &GmresOptions::default())?;
25//!
26//! assert!(result.converged);
27//! assert_eq!(result.solution.shape(), (1, 1));
28//! # Ok(())
29//! # }
30//! ```
31
32use std::collections::HashSet;
33
34use crate::any_scalar::AnyScalar;
35use crate::tensor_index::TensorIndex;
36use crate::tensor_like::{
37    DirectSumResult, FactorizeError, FactorizeOptions, FactorizeResult, LinearizationOrder,
38    TensorConstructionLike, TensorContractionLike, TensorFactorizationLike, TensorLike,
39    TensorVectorSpace, TensorVectorSpaceError,
40};
41
42/// A collection of tensors organized in a block structure.
43///
44/// Each block is a tensor of type `T` implementing [`TensorLike`].
45/// The flattened block list is ordered row-by-row:
46/// `(0, 0), (0, 1), ..., (1, 0), (1, 1), ...`.
47///
48/// # Type Parameters
49///
50/// * `T` - The tensor type for each block, must implement `TensorLike`
51#[derive(Debug, Clone)]
52pub struct BlockTensor<T: TensorLike> {
53    /// Blocks flattened row-by-row in block-matrix order
54    blocks: Vec<T>,
55    /// Block structure (rows, cols)
56    shape: (usize, usize),
57}
58
59impl<T: TensorLike> BlockTensor<T> {
60    /// Create a new block tensor with validation.
61    ///
62    /// # Arguments
63    ///
64    /// * `blocks` - Vector of blocks flattened row-by-row
65    /// * `shape` - Block structure as (rows, cols)
66    ///
67    /// # Errors
68    ///
69    /// Returns [`TensorVectorSpaceError`] when `rows * cols` overflows `usize`
70    /// or when `blocks.len()` differs from `rows * cols`.
71    ///
72    /// # Examples
73    ///
74    /// ```
75    /// use tensor4all_core::block_tensor::BlockTensor;
76    /// use tensor4all_core::{DynIndex, IdxTensor};
77    ///
78    /// let i = DynIndex::new_dyn(2);
79    /// let t1 = IdxTensor::from_dense(vec![i.clone()], vec![1.0, 2.0]).unwrap();
80    /// let t2 = IdxTensor::from_dense(vec![i.clone()], vec![3.0, 4.0]).unwrap();
81    ///
82    /// let bt = BlockTensor::try_new(vec![t1, t2], (2, 1)).unwrap();
83    /// assert_eq!(bt.shape(), (2, 1));
84    ///
85    /// // Wrong number of blocks returns an error
86    /// let t3 = IdxTensor::from_dense(vec![i], vec![5.0, 6.0]).unwrap();
87    /// assert!(BlockTensor::try_new(vec![t3], (2, 1)).is_err());
88    /// ```
89    pub fn try_new(
90        blocks: Vec<T>,
91        shape: (usize, usize),
92    ) -> std::result::Result<Self, TensorVectorSpaceError> {
93        let (rows, cols) = shape;
94        let expected = rows.checked_mul(cols).ok_or_else(|| {
95            TensorVectorSpaceError::from(anyhow::anyhow!(
96                "Block count shape ({rows}, {cols}) overflows usize"
97            ))
98        })?;
99        if expected != blocks.len() {
100            return Err(TensorVectorSpaceError::from(anyhow::anyhow!(
101                "Block count mismatch: shape ({}, {}) requires {} blocks, but got {}",
102                rows,
103                cols,
104                expected,
105                blocks.len()
106            )));
107        }
108        Ok(Self { blocks, shape })
109    }
110
111    /// Create a new block tensor.
112    ///
113    /// # Arguments
114    ///
115    /// * `blocks` - Vector of blocks flattened row-by-row
116    /// * `shape` - Block structure as (rows, cols)
117    ///
118    /// # Errors
119    ///
120    /// Returns [`TensorVectorSpaceError`] when `rows * cols` overflows `usize`
121    /// or when `blocks.len()` differs from `rows * cols`.
122    ///
123    /// # Examples
124    ///
125    /// ```
126    /// use tensor4all_core::block_tensor::BlockTensor;
127    /// use tensor4all_core::{DynIndex, IdxTensor};
128    ///
129    /// let i = DynIndex::new_dyn(2);
130    /// let t = IdxTensor::from_dense(vec![i], vec![1.0, 2.0]).unwrap();
131    /// let bt = BlockTensor::new(vec![t], (1, 1)).unwrap();
132    /// assert_eq!(bt.shape(), (1, 1));
133    /// ```
134    pub fn new(
135        blocks: Vec<T>,
136        shape: (usize, usize),
137    ) -> std::result::Result<Self, TensorVectorSpaceError> {
138        Self::try_new(blocks, shape)
139    }
140
141    /// Get the block structure (rows, cols).
142    ///
143    /// # Examples
144    ///
145    /// ```
146    /// use tensor4all_core::block_tensor::BlockTensor;
147    /// use tensor4all_core::{DynIndex, IdxTensor};
148    ///
149    /// let i = DynIndex::new_dyn(2);
150    /// let blocks: Vec<IdxTensor> = (0..6)
151    ///     .map(|_| IdxTensor::from_dense(vec![i.clone()], vec![0.0, 0.0]).unwrap())
152    ///     .collect();
153    /// let bt = BlockTensor::new(blocks, (2, 3)).unwrap();
154    /// assert_eq!(bt.shape(), (2, 3));
155    /// ```
156    pub fn shape(&self) -> (usize, usize) {
157        self.shape
158    }
159
160    /// Get the total number of blocks.
161    pub fn num_blocks(&self) -> usize {
162        self.blocks.len()
163    }
164
165    /// Get a reference to the block at (row, col).
166    ///
167    /// # Examples
168    ///
169    /// ```
170    /// use tensor4all_core::block_tensor::BlockTensor;
171    /// use tensor4all_core::{DynIndex, IdxTensor};
172    ///
173    /// let i = DynIndex::new_dyn(2);
174    /// let t1 = IdxTensor::from_dense(vec![i.clone()], vec![1.0, 2.0]).unwrap();
175    /// let t2 = IdxTensor::from_dense(vec![i], vec![3.0, 4.0]).unwrap();
176    /// let bt = BlockTensor::new(vec![t1, t2], (2, 1)).unwrap();
177    /// assert_eq!(bt.get(0, 0).unwrap().dims(), vec![2]);
178    /// assert_eq!(bt.get(1, 0).unwrap().dims(), vec![2]);
179    /// ```
180    pub fn get(&self, row: usize, col: usize) -> Option<&T> {
181        let (rows, cols) = self.shape;
182        if row >= rows || col >= cols {
183            return None;
184        }
185        self.blocks.get(row * cols + col)
186    }
187
188    /// Get a mutable reference to the block at (row, col).
189    ///
190    /// # Examples
191    ///
192    /// ```
193    /// use tensor4all_core::block_tensor::BlockTensor;
194    /// use tensor4all_core::{DynIndex, IdxTensor};
195    ///
196    /// let i = DynIndex::new_dyn(2);
197    /// let t = IdxTensor::from_dense(vec![i.clone()], vec![1.0, 2.0]).unwrap();
198    /// let mut bt = BlockTensor::new(vec![t], (1, 1)).unwrap();
199    /// let block = bt.get_mut(0, 0).unwrap();
200    /// assert_eq!(block.dims(), vec![2]);
201    /// ```
202    pub fn get_mut(&mut self, row: usize, col: usize) -> Option<&mut T> {
203        let (rows, cols) = self.shape;
204        if row >= rows || col >= cols {
205            return None;
206        }
207        self.blocks.get_mut(row * cols + col)
208    }
209
210    /// Get all blocks as a slice.
211    ///
212    /// # Examples
213    ///
214    /// ```
215    /// use tensor4all_core::block_tensor::BlockTensor;
216    /// use tensor4all_core::{DynIndex, IdxTensor};
217    ///
218    /// let i = DynIndex::new_dyn(2);
219    /// let t1 = IdxTensor::from_dense(vec![i.clone()], vec![1.0, 2.0]).unwrap();
220    /// let t2 = IdxTensor::from_dense(vec![i], vec![3.0, 4.0]).unwrap();
221    /// let bt = BlockTensor::new(vec![t1, t2], (1, 2)).unwrap();
222    /// assert_eq!(bt.blocks().len(), 2);
223    /// ```
224    pub fn blocks(&self) -> &[T] {
225        &self.blocks
226    }
227
228    /// Get all blocks as a mutable slice.
229    ///
230    /// # Examples
231    ///
232    /// ```
233    /// use tensor4all_core::block_tensor::BlockTensor;
234    /// use tensor4all_core::{DynIndex, IdxTensor};
235    ///
236    /// let i = DynIndex::new_dyn(2);
237    /// let t = IdxTensor::from_dense(vec![i], vec![1.0, 2.0]).unwrap();
238    /// let mut bt = BlockTensor::new(vec![t], (1, 1)).unwrap();
239    /// assert_eq!(bt.blocks_mut().len(), 1);
240    /// ```
241    pub fn blocks_mut(&mut self) -> &mut [T] {
242        &mut self.blocks
243    }
244
245    /// Consume self and return the blocks.
246    ///
247    /// # Examples
248    ///
249    /// ```
250    /// use tensor4all_core::block_tensor::BlockTensor;
251    /// use tensor4all_core::{DynIndex, IdxTensor};
252    ///
253    /// let i = DynIndex::new_dyn(2);
254    /// let t = IdxTensor::from_dense(vec![i], vec![1.0, 2.0]).unwrap();
255    /// let bt = BlockTensor::new(vec![t], (1, 1)).unwrap();
256    /// let blocks = bt.into_blocks();
257    /// assert_eq!(blocks.len(), 1);
258    /// assert_eq!(blocks[0].dims(), vec![2]);
259    /// ```
260    pub fn into_blocks(self) -> Vec<T> {
261        self.blocks
262    }
263
264    /// Validate that blocks share external indices consistently.
265    ///
266    /// For column vectors (cols=1), no index sharing is required between
267    /// blocks. Different rows can have independent physical indices
268    /// (the operator determines their relationship).
269    ///
270    /// For matrices (rows x cols), checks that:
271    /// - All blocks have the same number of external indices.
272    /// - Blocks in the same row share some common index IDs (output indices).
273    /// - Blocks in the same column share some common index IDs (input indices).
274    ///
275    /// # Errors
276    ///
277    /// Returns [`TensorVectorSpaceError`] when blocks disagree on their number
278    /// of external indices, or when blocks in a common row or column share no
279    /// index.
280    ///
281    /// # Examples
282    ///
283    /// ```
284    /// use tensor4all_core::block_tensor::BlockTensor;
285    /// use tensor4all_core::{DynIndex, IdxTensor};
286    ///
287    /// // Column vector: always valid regardless of index sharing
288    /// let i = DynIndex::new_dyn(2);
289    /// let j = DynIndex::new_dyn(2);
290    /// let t1 = IdxTensor::from_dense(vec![i], vec![1.0, 2.0]).unwrap();
291    /// let t2 = IdxTensor::from_dense(vec![j], vec![3.0, 4.0]).unwrap();
292    /// let bt = BlockTensor::new(vec![t1, t2], (2, 1)).unwrap();
293    /// assert!(bt.validate_indices().is_ok());
294    /// ```
295    pub fn validate_indices(&self) -> std::result::Result<(), TensorVectorSpaceError> {
296        let (rows, cols) = self.shape;
297
298        if cols <= 1 {
299            // Column vector: blocks in different rows can have independent indices.
300            // The operator determines the relationship between blocks.
301            return Ok(());
302        }
303
304        // Matrix: check all blocks have the same number of external indices
305        let first_count = self.blocks[0].num_external_indices();
306        for (i, block) in self.blocks.iter().enumerate().skip(1) {
307            let n = block.num_external_indices();
308            if n != first_count {
309                return Err(TensorVectorSpaceError::from(anyhow::anyhow!(
310                    "Block {} has {} external indices, but block 0 has {}",
311                    i,
312                    n,
313                    first_count
314                )));
315            }
316        }
317
318        // Same row: blocks should share at least one full output index.
319        for row in 0..rows {
320            let ref_indices: HashSet<_> = self
321                .get(row, 0)
322                .ok_or_else(|| anyhow::anyhow!("block index ({row}, 0) is out of bounds"))?
323                .external_indices()
324                .iter()
325                .cloned()
326                .collect();
327            for col in 1..cols {
328                let indices: HashSet<_> = self
329                    .get(row, col)
330                    .ok_or_else(|| anyhow::anyhow!("block index ({row}, {col}) is out of bounds"))?
331                    .external_indices()
332                    .iter()
333                    .cloned()
334                    .collect();
335                let common_count = ref_indices.intersection(&indices).count();
336                if common_count == 0 {
337                    return Err(TensorVectorSpaceError::from(anyhow::anyhow!(
338                        "Matrix row {}: blocks ({},{}) and ({},{}) share no common indices",
339                        row,
340                        row,
341                        0,
342                        row,
343                        col
344                    )));
345                }
346            }
347        }
348
349        // Same column: blocks should share at least one full input index.
350        for col in 0..cols {
351            let ref_indices: HashSet<_> = self
352                .get(0, col)
353                .ok_or_else(|| anyhow::anyhow!("block index (0, {col}) is out of bounds"))?
354                .external_indices()
355                .iter()
356                .cloned()
357                .collect();
358            for row in 1..rows {
359                let indices: HashSet<_> = self
360                    .get(row, col)
361                    .ok_or_else(|| anyhow::anyhow!("block index ({row}, {col}) is out of bounds"))?
362                    .external_indices()
363                    .iter()
364                    .cloned()
365                    .collect();
366                let common_count = ref_indices.intersection(&indices).count();
367                if common_count == 0 {
368                    return Err(TensorVectorSpaceError::from(anyhow::anyhow!(
369                        "Matrix col {}: blocks ({},{}) and ({},{}) share no common indices",
370                        col,
371                        0,
372                        col,
373                        row,
374                        col
375                    )));
376                }
377            }
378        }
379
380        Ok(())
381    }
382}
383
384// ============================================================================
385// TensorIndex implementation
386// ============================================================================
387
388impl<T: TensorLike> TensorIndex for BlockTensor<T> {
389    type Index = T::Index;
390    type Error = TensorVectorSpaceError;
391
392    fn external_indices(&self) -> Vec<Self::Index> {
393        // Collect unique external indices across all blocks (deduplicated by full index).
394        let mut seen = HashSet::new();
395        let mut result = Vec::new();
396        for block in &self.blocks {
397            for idx in block.external_indices() {
398                if seen.insert(idx.clone()) {
399                    result.push(idx);
400                }
401            }
402        }
403        result
404    }
405
406    fn replaceind(
407        &self,
408        old_index: &Self::Index,
409        new_index: &Self::Index,
410    ) -> std::result::Result<Self, Self::Error> {
411        let replaced: std::result::Result<Vec<T>, Self::Error> = self
412            .blocks
413            .iter()
414            .map(|b| {
415                b.replaceind(old_index, new_index)
416                    .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))
417            })
418            .collect();
419        Ok(Self {
420            blocks: replaced?,
421            shape: self.shape,
422        })
423    }
424
425    fn replace_indices(
426        &self,
427        old_indices: &[Self::Index],
428        new_indices: &[Self::Index],
429    ) -> std::result::Result<Self, Self::Error> {
430        let replaced: std::result::Result<Vec<T>, Self::Error> = self
431            .blocks
432            .iter()
433            .map(|b| {
434                b.replace_indices(old_indices, new_indices)
435                    .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))
436            })
437            .collect();
438        Ok(Self {
439            blocks: replaced?,
440            shape: self.shape,
441        })
442    }
443}
444
445// ============================================================================
446// TensorLike implementation
447// ============================================================================
448
449impl<T: TensorLike> TensorVectorSpace for BlockTensor<T> {
450    // ------------------------------------------------------------------------
451    // Vector space operations (required for GMRES)
452    // ------------------------------------------------------------------------
453
454    fn norm_squared(&self) -> std::result::Result<f64, Self::Error> {
455        self.blocks.iter().try_fold(0.0, |sum, block| {
456            let value = block
457                .norm_squared()
458                .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))?;
459            Ok(sum + value)
460        })
461    }
462
463    fn maxabs(&self) -> std::result::Result<f64, Self::Error> {
464        self.blocks.iter().try_fold(0.0_f64, |acc, block| {
465            let value = block
466                .maxabs()
467                .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))?;
468            Ok(acc.max(value))
469        })
470    }
471
472    fn scale(&self, scalar: AnyScalar) -> std::result::Result<Self, Self::Error> {
473        let scaled: std::result::Result<Vec<T>, Self::Error> = self
474            .blocks
475            .iter()
476            .map(|b| {
477                b.scale(scalar.clone())
478                    .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))
479            })
480            .collect();
481        Ok(Self {
482            blocks: scaled?,
483            shape: self.shape,
484        })
485    }
486
487    fn axpby(
488        &self,
489        a: AnyScalar,
490        other: &Self,
491        b: AnyScalar,
492    ) -> std::result::Result<Self, Self::Error> {
493        if self.shape != other.shape {
494            return Err(TensorVectorSpaceError::from(anyhow::anyhow!(
495                "Block shapes must match: {:?} vs {:?}",
496                self.shape,
497                other.shape
498            )));
499        }
500        let result: std::result::Result<Vec<T>, Self::Error> = self
501            .blocks
502            .iter()
503            .zip(other.blocks.iter())
504            .map(|(s, o)| {
505                s.axpby(a.clone(), o, b.clone())
506                    .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))
507            })
508            .collect();
509        Ok(Self {
510            blocks: result?,
511            shape: self.shape,
512        })
513    }
514
515    fn inner_product(&self, other: &Self) -> std::result::Result<AnyScalar, Self::Error> {
516        if self.shape != other.shape {
517            return Err(TensorVectorSpaceError::from(anyhow::anyhow!(
518                "Block shapes must match for inner product: {:?} vs {:?}",
519                self.shape,
520                other.shape
521            )));
522        }
523        let mut sum = AnyScalar::new_real(0.0);
524        for (s, o) in self.blocks.iter().zip(other.blocks.iter()) {
525            let value = s
526                .inner_product(o)
527                .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))?;
528            sum = sum + value;
529        }
530        Ok(sum)
531    }
532
533    fn validate(&self) -> std::result::Result<(), Self::Error> {
534        self.validate_indices()
535    }
536}
537
538impl<T: TensorLike> TensorContractionLike for BlockTensor<T> {
539    // ------------------------------------------------------------------------
540    // Tensor network operations
541    // ------------------------------------------------------------------------
542
543    fn conj(&self) -> Self {
544        let conjugated: Vec<T> = self.blocks.iter().map(|b| b.conj()).collect();
545        Self {
546            blocks: conjugated,
547            shape: self.shape,
548        }
549    }
550
551    fn direct_sum(
552        &self,
553        _other: &Self,
554        _pairs: &[(<Self as TensorIndex>::Index, <Self as TensorIndex>::Index)],
555    ) -> std::result::Result<DirectSumResult<Self>, Self::Error> {
556        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
557            "BlockTensor does not support direct_sum"
558        )))
559    }
560
561    fn outer_product(&self, _other: &Self) -> std::result::Result<Self, Self::Error> {
562        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
563            "BlockTensor does not support outer_product"
564        )))
565    }
566
567    fn permuteinds(
568        &self,
569        _new_order: &[<Self as TensorIndex>::Index],
570    ) -> std::result::Result<Self, Self::Error> {
571        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
572            "BlockTensor does not support permuteinds"
573        )))
574    }
575
576    fn fuse_indices(
577        &self,
578        old_indices: &[Self::Index],
579        new_index: Self::Index,
580        order: LinearizationOrder,
581    ) -> std::result::Result<Self, Self::Error> {
582        let blocks: std::result::Result<Vec<T>, Self::Error> = self
583            .blocks
584            .iter()
585            .map(|block| {
586                block
587                    .fuse_indices(old_indices, new_index.clone(), order)
588                    .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))
589            })
590            .collect();
591        Ok(Self {
592            blocks: blocks?,
593            shape: self.shape,
594        })
595    }
596
597    fn contract(_tensors: &[&Self]) -> std::result::Result<Self, Self::Error> {
598        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
599            "BlockTensor does not support contract"
600        )))
601    }
602}
603
604impl<T: TensorLike> TensorFactorizationLike for BlockTensor<T> {
605    fn factorize(
606        &self,
607        _left_inds: &[<Self as TensorIndex>::Index],
608        _options: &FactorizeOptions,
609    ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
610        Err(FactorizeError::ComputationError(anyhow::anyhow!(
611            "BlockTensor does not support factorize"
612        )))
613    }
614
615    fn factorize_full_rank(
616        &self,
617        _left_inds: &[<Self as TensorIndex>::Index],
618        _alg: crate::FactorizeAlg,
619        _canonical: crate::Canonical,
620    ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
621        Err(FactorizeError::ComputationError(anyhow::anyhow!(
622            "BlockTensor does not support factorize_full_rank"
623        )))
624    }
625}
626
627impl<T: TensorLike> TensorConstructionLike for BlockTensor<T> {
628    fn diagonal(
629        _input_index: &<Self as TensorIndex>::Index,
630        _output_index: &<Self as TensorIndex>::Index,
631    ) -> std::result::Result<Self, Self::Error> {
632        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
633            "BlockTensor does not support diagonal"
634        )))
635    }
636
637    fn scalar_one() -> std::result::Result<Self, Self::Error> {
638        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
639            "BlockTensor does not support scalar_one"
640        )))
641    }
642
643    fn ones(_indices: &[<Self as TensorIndex>::Index]) -> std::result::Result<Self, Self::Error> {
644        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
645            "BlockTensor does not support ones"
646        )))
647    }
648
649    fn onehot(
650        _index_vals: &[(<Self as TensorIndex>::Index, usize)],
651    ) -> std::result::Result<Self, Self::Error> {
652        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
653            "BlockTensor does not support onehot"
654        )))
655    }
656}
657
658#[cfg(test)]
659mod tests;
tensor4all_core/block_tensor.rs

tensor4all_core/
block_tensor.rs
