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 an error when the block count does not match the shape (a shape
70    /// /// mismatch).
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 an error when the block count does not match the shape (a shape
121    /// /// mismatch).
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 an error when the blocks have inconsistent indices (a shape or
278    /// /// index mismatch).
279    ///
280    /// # Examples
281    ///
282    /// ```
283    /// use tensor4all_core::block_tensor::BlockTensor;
284    /// use tensor4all_core::{DynIndex, IdxTensor};
285    ///
286    /// // Column vector: always valid regardless of index sharing
287    /// let i = DynIndex::new_dyn(2);
288    /// let j = DynIndex::new_dyn(2);
289    /// let t1 = IdxTensor::from_dense(vec![i], vec![1.0, 2.0]).unwrap();
290    /// let t2 = IdxTensor::from_dense(vec![j], vec![3.0, 4.0]).unwrap();
291    /// let bt = BlockTensor::new(vec![t1, t2], (2, 1)).unwrap();
292    /// assert!(bt.validate_indices().is_ok());
293    /// ```
294    pub fn validate_indices(&self) -> std::result::Result<(), TensorVectorSpaceError> {
295        let (rows, cols) = self.shape;
296
297        if cols <= 1 {
298            // Column vector: blocks in different rows can have independent indices.
299            // The operator determines the relationship between blocks.
300            return Ok(());
301        }
302
303        // Matrix: check all blocks have the same number of external indices
304        let first_count = self.blocks[0].num_external_indices();
305        for (i, block) in self.blocks.iter().enumerate().skip(1) {
306            let n = block.num_external_indices();
307            if n != first_count {
308                return Err(TensorVectorSpaceError::from(anyhow::anyhow!(
309                    "Block {} has {} external indices, but block 0 has {}",
310                    i,
311                    n,
312                    first_count
313                )));
314            }
315        }
316
317        // Same row: blocks should share at least one full output index.
318        for row in 0..rows {
319            let ref_indices: HashSet<_> = self
320                .get(row, 0)
321                .ok_or_else(|| anyhow::anyhow!("block index ({row}, 0) is out of bounds"))?
322                .external_indices()
323                .iter()
324                .cloned()
325                .collect();
326            for col in 1..cols {
327                let indices: HashSet<_> = self
328                    .get(row, col)
329                    .ok_or_else(|| anyhow::anyhow!("block index ({row}, {col}) is out of bounds"))?
330                    .external_indices()
331                    .iter()
332                    .cloned()
333                    .collect();
334                let common_count = ref_indices.intersection(&indices).count();
335                if common_count == 0 {
336                    return Err(TensorVectorSpaceError::from(anyhow::anyhow!(
337                        "Matrix row {}: blocks ({},{}) and ({},{}) share no common indices",
338                        row,
339                        row,
340                        0,
341                        row,
342                        col
343                    )));
344                }
345            }
346        }
347
348        // Same column: blocks should share at least one full input index.
349        for col in 0..cols {
350            let ref_indices: HashSet<_> = self
351                .get(0, col)
352                .ok_or_else(|| anyhow::anyhow!("block index (0, {col}) is out of bounds"))?
353                .external_indices()
354                .iter()
355                .cloned()
356                .collect();
357            for row in 1..rows {
358                let indices: HashSet<_> = self
359                    .get(row, col)
360                    .ok_or_else(|| anyhow::anyhow!("block index ({row}, {col}) is out of bounds"))?
361                    .external_indices()
362                    .iter()
363                    .cloned()
364                    .collect();
365                let common_count = ref_indices.intersection(&indices).count();
366                if common_count == 0 {
367                    return Err(TensorVectorSpaceError::from(anyhow::anyhow!(
368                        "Matrix col {}: blocks ({},{}) and ({},{}) share no common indices",
369                        col,
370                        0,
371                        col,
372                        row,
373                        col
374                    )));
375                }
376            }
377        }
378
379        Ok(())
380    }
381}
382
383// ============================================================================
384// TensorIndex implementation
385// ============================================================================
386
387impl<T: TensorLike> TensorIndex for BlockTensor<T> {
388    type Index = T::Index;
389    type Error = TensorVectorSpaceError;
390
391    fn external_indices(&self) -> Vec<Self::Index> {
392        // Collect unique external indices across all blocks (deduplicated by full index).
393        let mut seen = HashSet::new();
394        let mut result = Vec::new();
395        for block in &self.blocks {
396            for idx in block.external_indices() {
397                if seen.insert(idx.clone()) {
398                    result.push(idx);
399                }
400            }
401        }
402        result
403    }
404
405    fn replaceind(
406        &self,
407        old_index: &Self::Index,
408        new_index: &Self::Index,
409    ) -> std::result::Result<Self, Self::Error> {
410        let replaced: std::result::Result<Vec<T>, Self::Error> = self
411            .blocks
412            .iter()
413            .map(|b| {
414                b.replaceind(old_index, new_index)
415                    .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))
416            })
417            .collect();
418        Ok(Self {
419            blocks: replaced?,
420            shape: self.shape,
421        })
422    }
423
424    fn replace_indices(
425        &self,
426        old_indices: &[Self::Index],
427        new_indices: &[Self::Index],
428    ) -> std::result::Result<Self, Self::Error> {
429        let replaced: std::result::Result<Vec<T>, Self::Error> = self
430            .blocks
431            .iter()
432            .map(|b| {
433                b.replace_indices(old_indices, new_indices)
434                    .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))
435            })
436            .collect();
437        Ok(Self {
438            blocks: replaced?,
439            shape: self.shape,
440        })
441    }
442}
443
444// ============================================================================
445// TensorLike implementation
446// ============================================================================
447
448impl<T: TensorLike> TensorVectorSpace for BlockTensor<T> {
449    // ------------------------------------------------------------------------
450    // Vector space operations (required for GMRES)
451    // ------------------------------------------------------------------------
452
453    fn norm_squared(&self) -> std::result::Result<f64, Self::Error> {
454        self.blocks.iter().try_fold(0.0, |sum, block| {
455            let value = block
456                .norm_squared()
457                .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))?;
458            Ok(sum + value)
459        })
460    }
461
462    fn maxabs(&self) -> std::result::Result<f64, Self::Error> {
463        self.blocks.iter().try_fold(0.0_f64, |acc, block| {
464            let value = block
465                .maxabs()
466                .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))?;
467            Ok(acc.max(value))
468        })
469    }
470
471    fn scale(&self, scalar: AnyScalar) -> std::result::Result<Self, Self::Error> {
472        let scaled: std::result::Result<Vec<T>, Self::Error> = self
473            .blocks
474            .iter()
475            .map(|b| {
476                b.scale(scalar.clone())
477                    .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))
478            })
479            .collect();
480        Ok(Self {
481            blocks: scaled?,
482            shape: self.shape,
483        })
484    }
485
486    fn axpby(
487        &self,
488        a: AnyScalar,
489        other: &Self,
490        b: AnyScalar,
491    ) -> std::result::Result<Self, Self::Error> {
492        if self.shape != other.shape {
493            return Err(TensorVectorSpaceError::from(anyhow::anyhow!(
494                "Block shapes must match: {:?} vs {:?}",
495                self.shape,
496                other.shape
497            )));
498        }
499        let result: std::result::Result<Vec<T>, Self::Error> = self
500            .blocks
501            .iter()
502            .zip(other.blocks.iter())
503            .map(|(s, o)| {
504                s.axpby(a.clone(), o, b.clone())
505                    .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))
506            })
507            .collect();
508        Ok(Self {
509            blocks: result?,
510            shape: self.shape,
511        })
512    }
513
514    fn inner_product(&self, other: &Self) -> std::result::Result<AnyScalar, Self::Error> {
515        if self.shape != other.shape {
516            return Err(TensorVectorSpaceError::from(anyhow::anyhow!(
517                "Block shapes must match for inner product: {:?} vs {:?}",
518                self.shape,
519                other.shape
520            )));
521        }
522        let mut sum = AnyScalar::new_real(0.0);
523        for (s, o) in self.blocks.iter().zip(other.blocks.iter()) {
524            let value = s
525                .inner_product(o)
526                .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))?;
527            sum = sum + value;
528        }
529        Ok(sum)
530    }
531
532    fn validate(&self) -> std::result::Result<(), Self::Error> {
533        self.validate_indices()
534    }
535}
536
537impl<T: TensorLike> TensorContractionLike for BlockTensor<T> {
538    // ------------------------------------------------------------------------
539    // Tensor network operations
540    // ------------------------------------------------------------------------
541
542    fn conj(&self) -> Self {
543        let conjugated: Vec<T> = self.blocks.iter().map(|b| b.conj()).collect();
544        Self {
545            blocks: conjugated,
546            shape: self.shape,
547        }
548    }
549
550    fn direct_sum(
551        &self,
552        _other: &Self,
553        _pairs: &[(<Self as TensorIndex>::Index, <Self as TensorIndex>::Index)],
554    ) -> std::result::Result<DirectSumResult<Self>, Self::Error> {
555        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
556            "BlockTensor does not support direct_sum"
557        )))
558    }
559
560    fn outer_product(&self, _other: &Self) -> std::result::Result<Self, Self::Error> {
561        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
562            "BlockTensor does not support outer_product"
563        )))
564    }
565
566    fn permuteinds(
567        &self,
568        _new_order: &[<Self as TensorIndex>::Index],
569    ) -> std::result::Result<Self, Self::Error> {
570        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
571            "BlockTensor does not support permuteinds"
572        )))
573    }
574
575    fn fuse_indices(
576        &self,
577        old_indices: &[Self::Index],
578        new_index: Self::Index,
579        order: LinearizationOrder,
580    ) -> std::result::Result<Self, Self::Error> {
581        let blocks: std::result::Result<Vec<T>, Self::Error> = self
582            .blocks
583            .iter()
584            .map(|block| {
585                block
586                    .fuse_indices(old_indices, new_index.clone(), order)
587                    .map_err(|error| TensorVectorSpaceError::from(anyhow::Error::new(error)))
588            })
589            .collect();
590        Ok(Self {
591            blocks: blocks?,
592            shape: self.shape,
593        })
594    }
595
596    fn contract(_tensors: &[&Self]) -> std::result::Result<Self, Self::Error> {
597        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
598            "BlockTensor does not support contract"
599        )))
600    }
601}
602
603impl<T: TensorLike> TensorFactorizationLike for BlockTensor<T> {
604    fn factorize(
605        &self,
606        _left_inds: &[<Self as TensorIndex>::Index],
607        _options: &FactorizeOptions,
608    ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
609        Err(FactorizeError::ComputationError(anyhow::anyhow!(
610            "BlockTensor does not support factorize"
611        )))
612    }
613
614    fn factorize_full_rank(
615        &self,
616        _left_inds: &[<Self as TensorIndex>::Index],
617        _alg: crate::FactorizeAlg,
618        _canonical: crate::Canonical,
619    ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
620        Err(FactorizeError::ComputationError(anyhow::anyhow!(
621            "BlockTensor does not support factorize_full_rank"
622        )))
623    }
624}
625
626impl<T: TensorLike> TensorConstructionLike for BlockTensor<T> {
627    fn diagonal(
628        _input_index: &<Self as TensorIndex>::Index,
629        _output_index: &<Self as TensorIndex>::Index,
630    ) -> std::result::Result<Self, Self::Error> {
631        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
632            "BlockTensor does not support diagonal"
633        )))
634    }
635
636    fn scalar_one() -> std::result::Result<Self, Self::Error> {
637        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
638            "BlockTensor does not support scalar_one"
639        )))
640    }
641
642    fn ones(_indices: &[<Self as TensorIndex>::Index]) -> std::result::Result<Self, Self::Error> {
643        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
644            "BlockTensor does not support ones"
645        )))
646    }
647
648    fn onehot(
649        _index_vals: &[(<Self as TensorIndex>::Index, usize)],
650    ) -> std::result::Result<Self, Self::Error> {
651        Err(TensorVectorSpaceError::from(anyhow::anyhow!(
652            "BlockTensor does not support onehot"
653        )))
654    }
655}
656
657#[cfg(test)]
658mod tests;
tensor4all_core/block_tensor.rs

tensor4all_core/
block_tensor.rs
