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 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
196pub 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#[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 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
408pub 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#[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}