Skip to main content

qmbed/
block.rs

1use std::sync::Arc;
2
3use num_complex::Complex64;
4
5use crate::operator::{
6    LinearOperator, MatrixFormat, Operator, TimeDependentOperator, TimeOperator, check_apply_shape,
7};
8use crate::{QmbedError, Result};
9
10fn block_offsets(shapes: impl IntoIterator<Item = (usize, usize)>) -> Result<Vec<usize>> {
11    let mut offsets = vec![0_usize];
12    for shape in shapes {
13        if shape.0 != shape.1 {
14            return Err(QmbedError::DimensionMismatch(
15                "block-diagonal operators require square blocks".into(),
16            ));
17        }
18        offsets.push(
19            offsets
20                .last()
21                .copied()
22                .unwrap_or_default()
23                .checked_add(shape.0)
24                .ok_or_else(|| QmbedError::UnsupportedBackend("block dimension overflow".into()))?,
25        );
26    }
27    if offsets.len() == 1 {
28        return Err(QmbedError::InvalidOptions(
29            "at least one operator block is required".into(),
30        ));
31    }
32    Ok(offsets)
33}
34
35#[derive(Clone)]
36struct ProjectionLayout {
37    projectors: Vec<Arc<dyn LinearOperator>>,
38    offsets: Vec<usize>,
39    full_dimension: usize,
40}
41
42impl ProjectionLayout {
43    fn new(
44        block_shapes: impl IntoIterator<Item = (usize, usize)>,
45        projectors: impl IntoIterator<Item = Arc<dyn LinearOperator>>,
46        tolerance: f64,
47    ) -> Result<Self> {
48        if !tolerance.is_finite() || tolerance <= 0.0 {
49            return Err(QmbedError::InvalidOptions(
50                "projector-isometry tolerance must be positive and finite".into(),
51            ));
52        }
53        let block_shapes: Vec<_> = block_shapes.into_iter().collect();
54        let offsets = block_offsets(block_shapes.iter().copied())?;
55        let projectors: Vec<_> = projectors.into_iter().collect();
56        if projectors.len() != block_shapes.len() {
57            return Err(QmbedError::DimensionMismatch(format!(
58                "{} blocks require the same number of projectors, got {}",
59                block_shapes.len(),
60                projectors.len()
61            )));
62        }
63        let full_dimension = projectors
64            .first()
65            .map(|projector| projector.shape().0)
66            .unwrap_or_default();
67        if full_dimension == 0 {
68            return Err(QmbedError::DimensionMismatch(
69                "projectors must have a positive parent-space dimension".into(),
70            ));
71        }
72        for (index, (block_shape, projector)) in
73            block_shapes.iter().zip(projectors.iter()).enumerate()
74        {
75            let projector_shape = projector.shape();
76            if projector_shape != (full_dimension, block_shape.0) {
77                return Err(QmbedError::DimensionMismatch(format!(
78                    "projector {index} has shape {projector_shape:?}, expected ({full_dimension}, {})",
79                    block_shape.0
80                )));
81            }
82        }
83
84        // Validate the combined map P = [P_0 ... P_n] as one isometry.
85        // This checks both normalization inside a block and mutual
86        // orthogonality between different sectors without requiring stored
87        // projector matrices.
88        for (source_index, source) in projectors.iter().enumerate() {
89            let source_dimension = block_shapes[source_index].0;
90            let mut coordinate = vec![Complex64::new(0.0, 0.0); source_dimension];
91            let mut parent = vec![Complex64::new(0.0, 0.0); full_dimension];
92            for source_column in 0..source_dimension {
93                coordinate.fill(Complex64::new(0.0, 0.0));
94                coordinate[source_column] = Complex64::new(1.0, 0.0);
95                source.apply(&coordinate, &mut parent)?;
96                for (target_index, target) in projectors.iter().enumerate() {
97                    let target_dimension = block_shapes[target_index].0;
98                    let mut overlap = vec![Complex64::new(0.0, 0.0); target_dimension];
99                    target.apply_adjoint(&parent, &mut overlap)?;
100                    for (target_column, value) in overlap.into_iter().enumerate() {
101                        let expected =
102                            if source_index == target_index && source_column == target_column {
103                                Complex64::new(1.0, 0.0)
104                            } else {
105                                Complex64::new(0.0, 0.0)
106                            };
107                        if (value - expected).norm() > tolerance {
108                            return Err(QmbedError::InvalidOptions(format!(
109                                "combined block projector is not an isometry within tolerance {tolerance}: \
110                                 block {target_index} column {target_column} overlaps block \
111                                 {source_index} column {source_column} by {value}"
112                            )));
113                        }
114                    }
115                }
116            }
117        }
118
119        Ok(Self {
120            projectors,
121            offsets,
122            full_dimension,
123        })
124    }
125
126    fn block_dimension(&self) -> usize {
127        self.offsets.last().copied().unwrap_or_default()
128    }
129
130    fn project(&self, parent: &[Complex64], blocks: &mut [Complex64]) -> Result<()> {
131        if parent.len() != self.full_dimension || blocks.len() != self.block_dimension() {
132            return Err(QmbedError::DimensionMismatch(format!(
133                "projecting from parent dimension {} requires input length {} and block output length {}, got {} and {}",
134                self.full_dimension,
135                self.full_dimension,
136                self.block_dimension(),
137                parent.len(),
138                blocks.len()
139            )));
140        }
141        for (index, projector) in self.projectors.iter().enumerate() {
142            let start = self.offsets[index];
143            let end = self.offsets[index + 1];
144            projector.apply_adjoint(parent, &mut blocks[start..end])?;
145        }
146        Ok(())
147    }
148
149    fn lift(&self, blocks: &[Complex64], parent: &mut [Complex64]) -> Result<()> {
150        if blocks.len() != self.block_dimension() || parent.len() != self.full_dimension {
151            return Err(QmbedError::DimensionMismatch(format!(
152                "lifting from block dimension {} requires input length {} and parent output length {}, got {} and {}",
153                self.block_dimension(),
154                self.block_dimension(),
155                self.full_dimension,
156                blocks.len(),
157                parent.len()
158            )));
159        }
160        parent.fill(Complex64::new(0.0, 0.0));
161        let mut contribution = vec![Complex64::new(0.0, 0.0); self.full_dimension];
162        for (index, projector) in self.projectors.iter().enumerate() {
163            let start = self.offsets[index];
164            let end = self.offsets[index + 1];
165            projector.apply(&blocks[start..end], &mut contribution)?;
166            for (value, added) in parent.iter_mut().zip(contribution.iter()) {
167                *value += *added;
168            }
169        }
170        Ok(())
171    }
172
173    fn completeness_residual(&self) -> Result<f64> {
174        let mut parent = vec![Complex64::new(0.0, 0.0); self.full_dimension];
175        let mut blocks = vec![Complex64::new(0.0, 0.0); self.block_dimension()];
176        let mut reconstructed = vec![Complex64::new(0.0, 0.0); self.full_dimension];
177        let mut maximum = 0.0_f64;
178        for column in 0..self.full_dimension {
179            parent.fill(Complex64::new(0.0, 0.0));
180            parent[column] = Complex64::new(1.0, 0.0);
181            self.project(&parent, &mut blocks)?;
182            self.lift(&blocks, &mut reconstructed)?;
183            for (row, value) in reconstructed.iter().copied().enumerate() {
184                let expected = if row == column {
185                    Complex64::new(1.0, 0.0)
186                } else {
187                    Complex64::new(0.0, 0.0)
188                };
189                maximum = maximum.max((value - expected).norm());
190            }
191        }
192        Ok(maximum)
193    }
194}
195
196/// Delayed matrix-free direct sum of static blocks.
197pub struct BlockOps {
198    blocks: Vec<Arc<dyn LinearOperator>>,
199    offsets: Vec<usize>,
200}
201
202impl BlockOps {
203    pub fn new(blocks: impl IntoIterator<Item = Arc<dyn LinearOperator>>) -> Result<Self> {
204        let blocks: Vec<_> = blocks.into_iter().collect();
205        let offsets = block_offsets(blocks.iter().map(|block| block.shape()))?;
206        Ok(Self { blocks, offsets })
207    }
208
209    pub fn blocks(&self) -> usize {
210        self.blocks.len()
211    }
212
213    pub fn materialize(&self, format: MatrixFormat) -> Result<Operator> {
214        block_diag_hamiltonian(self.blocks.iter().cloned(), format)
215    }
216
217    pub fn push(&mut self, block: Arc<dyn LinearOperator>) -> Result<()> {
218        if block.shape().0 != block.shape().1 {
219            return Err(QmbedError::DimensionMismatch(
220                "block-diagonal operators require square blocks".into(),
221            ));
222        }
223        let next = self
224            .offsets
225            .last()
226            .copied()
227            .unwrap_or_default()
228            .checked_add(block.shape().0)
229            .ok_or_else(|| QmbedError::UnsupportedBackend("block dimension overflow".into()))?;
230        self.blocks.push(block);
231        self.offsets.push(next);
232        Ok(())
233    }
234}
235
236impl LinearOperator for BlockOps {
237    fn shape(&self) -> (usize, usize) {
238        let dimension = self.offsets.last().copied().unwrap_or_default();
239        (dimension, dimension)
240    }
241
242    fn format(&self) -> MatrixFormat {
243        MatrixFormat::MatrixFree
244    }
245
246    fn apply(&self, input: &[Complex64], output: &mut [Complex64]) -> Result<()> {
247        check_apply_shape(self.shape(), input, output)?;
248        output.fill(Complex64::new(0.0, 0.0));
249        for (index, block) in self.blocks.iter().enumerate() {
250            let start = self.offsets[index];
251            let end = self.offsets[index + 1];
252            block.apply(&input[start..end], &mut output[start..end])?;
253        }
254        Ok(())
255    }
256
257    fn stored_triplets(&self) -> Result<Option<Vec<(usize, usize, Complex64)>>> {
258        let mut entries = Vec::new();
259        for (block_index, block) in self.blocks.iter().enumerate() {
260            let Some(block_entries) = block.stored_triplets()? else {
261                return Ok(None);
262            };
263            let offset = self.offsets[block_index];
264            entries.extend(
265                block_entries
266                    .into_iter()
267                    .map(|(row, column, value)| (offset + row, offset + column, value)),
268            );
269        }
270        Ok(Some(entries))
271    }
272}
273
274/// A direct sum of sector operators acting in a shared parent Hilbert space.
275///
276/// Each projector maps one sector coordinate vector into the parent space.
277/// The combined projector is validated as an isometry, so this operator
278/// represents `P (⊕ H_sector) P†` without materializing either the direct sum
279/// or the parent-space matrix.
280#[derive(Clone)]
281pub struct ProjectedBlockOps {
282    blocks: Vec<Arc<dyn LinearOperator>>,
283    projection: ProjectionLayout,
284}
285
286impl ProjectedBlockOps {
287    pub fn new(
288        blocks: impl IntoIterator<Item = Arc<dyn LinearOperator>>,
289        projectors: impl IntoIterator<Item = Arc<dyn LinearOperator>>,
290        tolerance: f64,
291    ) -> Result<Self> {
292        let blocks: Vec<_> = blocks.into_iter().collect();
293        let projection = ProjectionLayout::new(
294            blocks.iter().map(|block| block.shape()),
295            projectors,
296            tolerance,
297        )?;
298        Ok(Self { blocks, projection })
299    }
300
301    pub fn blocks(&self) -> usize {
302        self.blocks.len()
303    }
304
305    pub fn full_dimension(&self) -> usize {
306        self.projection.full_dimension
307    }
308
309    pub fn block_dimension(&self) -> usize {
310        self.projection.block_dimension()
311    }
312
313    pub fn project(&self, parent: &[Complex64], blocks: &mut [Complex64]) -> Result<()> {
314        self.projection.project(parent, blocks)
315    }
316
317    pub fn lift(&self, blocks: &[Complex64], parent: &mut [Complex64]) -> Result<()> {
318        self.projection.lift(blocks, parent)
319    }
320
321    /// Maximum entrywise residual of `P P† - I` in the parent basis.
322    ///
323    /// An isometric projector collection need not be complete: selected
324    /// sectors can intentionally span only a subspace. Callers which require a
325    /// full decomposition can enforce their own tolerance on this residual.
326    pub fn completeness_residual(&self) -> Result<f64> {
327        self.projection.completeness_residual()
328    }
329
330    pub fn materialize(&self, format: MatrixFormat) -> Result<Operator> {
331        Operator::from_triplets(
332            self.shape().0,
333            self.shape().1,
334            streamed_triplets(self)?,
335            format,
336        )
337    }
338}
339
340impl LinearOperator for ProjectedBlockOps {
341    fn shape(&self) -> (usize, usize) {
342        (
343            self.projection.full_dimension,
344            self.projection.full_dimension,
345        )
346    }
347
348    fn format(&self) -> MatrixFormat {
349        MatrixFormat::MatrixFree
350    }
351
352    fn apply(&self, input: &[Complex64], output: &mut [Complex64]) -> Result<()> {
353        check_apply_shape(self.shape(), input, output)?;
354        let mut block_input = vec![Complex64::new(0.0, 0.0); self.projection.block_dimension()];
355        let mut block_output = block_input.clone();
356        self.projection.project(input, &mut block_input)?;
357        for (index, block) in self.blocks.iter().enumerate() {
358            let start = self.projection.offsets[index];
359            let end = self.projection.offsets[index + 1];
360            block.apply(&block_input[start..end], &mut block_output[start..end])?;
361        }
362        self.projection.lift(&block_output, output)
363    }
364}
365
366fn streamed_triplets(
367    operator: &(impl LinearOperator + ?Sized),
368) -> Result<Vec<(usize, usize, Complex64)>> {
369    if let Some(entries) = operator.stored_triplets()? {
370        return Ok(entries);
371    }
372    let shape = operator.shape();
373    let mut input = vec![Complex64::new(0.0, 0.0); shape.1];
374    let mut output = vec![Complex64::new(0.0, 0.0); shape.0];
375    let mut entries = Vec::new();
376    for column in 0..shape.1 {
377        input.fill(Complex64::new(0.0, 0.0));
378        input[column] = Complex64::new(1.0, 0.0);
379        operator.apply(&input, &mut output)?;
380        for (row, value) in output.iter().copied().enumerate() {
381            if value.norm() > f64::EPSILON {
382                entries.push((row, column, value));
383            }
384        }
385    }
386    Ok(entries)
387}
388
389pub fn block_diag_hamiltonian(
390    blocks: impl IntoIterator<Item = Arc<dyn LinearOperator>>,
391    format: MatrixFormat,
392) -> Result<Operator> {
393    let blocks: Vec<_> = blocks.into_iter().collect();
394    let offsets = block_offsets(blocks.iter().map(|block| block.shape()))?;
395    let dimension = offsets.last().copied().unwrap_or_default();
396    let mut triplets = Vec::new();
397    for (block_index, block) in blocks.iter().enumerate() {
398        let offset = offsets[block_index];
399        triplets.extend(
400            streamed_triplets(block.as_ref())?
401                .into_iter()
402                .map(|(row, column, value)| (offset + row, offset + column, value)),
403        );
404    }
405    Operator::from_triplets(dimension, dimension, triplets, format)
406}
407
408/// Delayed direct sum of explicitly time-dependent blocks.
409pub struct DynamicBlockOps {
410    blocks: Vec<Arc<dyn TimeDependentOperator>>,
411    offsets: Vec<usize>,
412}
413
414impl DynamicBlockOps {
415    pub fn new(blocks: impl IntoIterator<Item = Arc<dyn TimeDependentOperator>>) -> Result<Self> {
416        let blocks: Vec<_> = blocks.into_iter().collect();
417        let offsets = block_offsets(blocks.iter().map(|block| block.shape()))?;
418        Ok(Self { blocks, offsets })
419    }
420
421    pub fn blocks(&self) -> usize {
422        self.blocks.len()
423    }
424
425    pub fn push(&mut self, block: Arc<dyn TimeDependentOperator>) -> Result<()> {
426        if block.shape().0 != block.shape().1 {
427            return Err(QmbedError::DimensionMismatch(
428                "block-diagonal operators require square blocks".into(),
429            ));
430        }
431        let next = self
432            .offsets
433            .last()
434            .copied()
435            .unwrap_or_default()
436            .checked_add(block.shape().0)
437            .ok_or_else(|| QmbedError::UnsupportedBackend("block dimension overflow".into()))?;
438        self.blocks.push(block);
439        self.offsets.push(next);
440        Ok(())
441    }
442
443    pub fn materialize(&self, time: f64, format: MatrixFormat) -> Result<Operator> {
444        if !time.is_finite() {
445            return Err(QmbedError::InvalidOptions(
446                "dynamic block materialization time must be finite".into(),
447            ));
448        }
449        let dimension = self.offsets.last().copied().unwrap_or_default();
450        let mut entries = Vec::new();
451        for (block_index, block) in self.blocks.iter().enumerate() {
452            let block_dimension = block.shape().0;
453            let offset = self.offsets[block_index];
454            let mut input = vec![Complex64::new(0.0, 0.0); block_dimension];
455            let mut output = vec![Complex64::new(0.0, 0.0); block_dimension];
456            for column in 0..block_dimension {
457                input.fill(Complex64::new(0.0, 0.0));
458                input[column] = Complex64::new(1.0, 0.0);
459                block.apply_at(time, &input, &mut output)?;
460                for (row, value) in output.iter().copied().enumerate() {
461                    if value.norm() > f64::EPSILON {
462                        entries.push((offset + row, offset + column, value));
463                    }
464                }
465            }
466        }
467        Operator::from_triplets(dimension, dimension, entries, format)
468    }
469}
470
471impl TimeDependentOperator for DynamicBlockOps {
472    fn shape(&self) -> (usize, usize) {
473        let dimension = self.offsets.last().copied().unwrap_or_default();
474        (dimension, dimension)
475    }
476
477    fn apply_at(&self, time: f64, input: &[Complex64], output: &mut [Complex64]) -> Result<()> {
478        check_apply_shape(self.shape(), input, output)?;
479        output.fill(Complex64::new(0.0, 0.0));
480        for (index, block) in self.blocks.iter().enumerate() {
481            let start = self.offsets[index];
482            let end = self.offsets[index + 1];
483            block.apply_at(time, &input[start..end], &mut output[start..end])?;
484        }
485        Ok(())
486    }
487}
488
489/// Time-dependent sector operators lifted into a shared parent Hilbert space.
490///
491/// This is the dynamic analogue of [`ProjectedBlockOps`] and evaluates
492/// `P (⊕ H_sector(t)) P†` matrix-free at every requested time.
493#[derive(Clone)]
494pub struct ProjectedDynamicBlockOps {
495    blocks: Vec<Arc<dyn TimeDependentOperator>>,
496    projection: ProjectionLayout,
497}
498
499impl ProjectedDynamicBlockOps {
500    pub fn new(
501        blocks: impl IntoIterator<Item = Arc<dyn TimeDependentOperator>>,
502        projectors: impl IntoIterator<Item = Arc<dyn LinearOperator>>,
503        tolerance: f64,
504    ) -> Result<Self> {
505        let blocks: Vec<_> = blocks.into_iter().collect();
506        let projection = ProjectionLayout::new(
507            blocks.iter().map(|block| block.shape()),
508            projectors,
509            tolerance,
510        )?;
511        Ok(Self { blocks, projection })
512    }
513
514    pub fn blocks(&self) -> usize {
515        self.blocks.len()
516    }
517
518    pub fn full_dimension(&self) -> usize {
519        self.projection.full_dimension
520    }
521
522    pub fn block_dimension(&self) -> usize {
523        self.projection.block_dimension()
524    }
525
526    pub fn project(&self, parent: &[Complex64], blocks: &mut [Complex64]) -> Result<()> {
527        self.projection.project(parent, blocks)
528    }
529
530    pub fn lift(&self, blocks: &[Complex64], parent: &mut [Complex64]) -> Result<()> {
531        self.projection.lift(blocks, parent)
532    }
533
534    pub fn completeness_residual(&self) -> Result<f64> {
535        self.projection.completeness_residual()
536    }
537
538    pub fn materialize(&self, time: f64, format: MatrixFormat) -> Result<Operator> {
539        TimeOperator::from_operator(Arc::new(self.clone())).evaluate(time, format)
540    }
541}
542
543impl TimeDependentOperator for ProjectedDynamicBlockOps {
544    fn shape(&self) -> (usize, usize) {
545        (
546            self.projection.full_dimension,
547            self.projection.full_dimension,
548        )
549    }
550
551    fn apply_at(&self, time: f64, input: &[Complex64], output: &mut [Complex64]) -> Result<()> {
552        if !time.is_finite() {
553            return Err(QmbedError::InvalidOptions(
554                "dynamic block evaluation time must be finite".into(),
555            ));
556        }
557        check_apply_shape(self.shape(), input, output)?;
558        let mut block_input = vec![Complex64::new(0.0, 0.0); self.projection.block_dimension()];
559        let mut block_output = block_input.clone();
560        self.projection.project(input, &mut block_input)?;
561        for (index, block) in self.blocks.iter().enumerate() {
562            let start = self.projection.offsets[index];
563            let end = self.projection.offsets[index + 1];
564            block.apply_at(
565                time,
566                &block_input[start..end],
567                &mut block_output[start..end],
568            )?;
569        }
570        self.projection.lift(&block_output, output)
571    }
572}