treams_core/cluster/
interaction.rs

1//! Dense solves of `(I - T C) X = T B` through one LU factorization of `I - T C`.
2//!
3//! [`interaction`] returns the full interacting T-matrix `X = (I - T C)⁻¹ T`.
4//! [`InteractionFactor`] keeps the factorization and solves only requested incident
5//! fields `B`, any number of times.
6//!
7//! Upstream: `treams.TMatrix.interaction.solve()` for [`interaction`].
8//! [`InteractionFactor`] is a treams-rs extension.
9
10use std::sync::Arc;
11
12#[path = "interaction_state.rs"]
13mod state;
14
15use nalgebra::DMatrix;
16
17use crate::{
18    Complex, Error, Result,
19    linalg::{
20        Lu, product, product_adjoint_right, product_adjoint_right_into, product_into,
21        product_views, view, view_mut,
22    },
23    numerics::{self, finite},
24};
25
26/// Block-diagonal local T-matrix; a dense local matrix is a single block.
27#[derive(Clone, Debug)]
28pub(crate) struct LocalMatrix {
29    blocks: Vec<DMatrix<Complex>>,
30}
31
32impl LocalMatrix {
33    /// One dense local matrix.
34    pub(crate) fn from_dense(local: DMatrix<Complex>) -> Result<Self> {
35        if !local.is_square() || local.nrows() == 0 || local.iter().any(|&z| !finite(z)) {
36            return Err(Error::InvalidInput(
37                "local matrix must be finite, square and nonempty".into(),
38            ));
39        }
40        Ok(Self {
41            blocks: vec![local],
42        })
43    }
44
45    /// The diagonal blocks of a block-diagonal local matrix, one per particle.
46    pub(crate) fn from_blocks(blocks: Vec<DMatrix<Complex>>) -> Result<Self> {
47        if blocks.is_empty()
48            || blocks
49                .iter()
50                .any(|m| !m.is_square() || m.nrows() == 0 || m.iter().any(|z| !finite(*z)))
51        {
52            return Err(Error::InvalidInput(
53                "local blocks must be finite nonempty square matrices".into(),
54            ));
55        }
56        Ok(Self { blocks })
57    }
58
59    pub(crate) fn dimension(&self) -> usize {
60        self.blocks.iter().map(DMatrix::nrows).sum()
61    }
62
63    /// The dense matrix, with zeros off the diagonal blocks.
64    fn to_dense(&self) -> Result<DMatrix<Complex>> {
65        let mut result = numerics::zeros(self.dimension(), self.dimension())?;
66        let mut offset = 0;
67        for block in &self.blocks {
68            result
69                .view_mut((offset, offset), block.shape())
70                .copy_from(block);
71            offset += block.nrows();
72        }
73        Ok(result)
74    }
75
76    /// `T right`, or `Tᴴ right` if `adjoint`.
77    pub(crate) fn apply(
78        &self,
79        right: &DMatrix<Complex>,
80        adjoint: bool,
81    ) -> Result<DMatrix<Complex>> {
82        let mut result = numerics::zeros(self.dimension(), right.ncols())?;
83        self.apply_into(&mut result, right, adjoint);
84        Ok(result)
85    }
86
87    fn apply_into(&self, result: &mut DMatrix<Complex>, right: &DMatrix<Complex>, adjoint: bool) {
88        Self::apply_blocks_into(&self.blocks, result, right, adjoint);
89    }
90
91    fn apply_blocks_into(
92        blocks: &[DMatrix<Complex>],
93        result: &mut DMatrix<Complex>,
94        right: &DMatrix<Complex>,
95        adjoint: bool,
96    ) {
97        let mut offset = 0;
98        for block in blocks {
99            let target = view_mut(result).subrows_mut(offset, block.nrows());
100            let rhs = view(right).subrows(offset, block.nrows());
101            if adjoint {
102                product_into(target, view(block).adjoint(), rhs);
103            } else {
104                product_into(target, view(block), rhs);
105            }
106            offset += block.nrows();
107        }
108    }
109
110    fn validate_tangent(&self, tangent: &[DMatrix<Complex>]) -> Result<()> {
111        if tangent.len() != self.blocks.len()
112            || tangent.iter().zip(&self.blocks).any(|(direction, block)| {
113                direction.shape() != block.shape() || direction.iter().any(|&z| !finite(z))
114            })
115        {
116            return Err(Error::InvalidInput(
117                "local tangent blocks must be finite and match the local matrices".into(),
118            ));
119        }
120        Ok(())
121    }
122
123    /// Apply a validated block tangent without materializing its off-diagonal zeros.
124    fn apply_tangent(
125        &self,
126        tangent: &[DMatrix<Complex>],
127        right: &DMatrix<Complex>,
128    ) -> Result<DMatrix<Complex>> {
129        let mut result = numerics::zeros(self.dimension(), right.ncols())?;
130        Self::apply_blocks_into(tangent, &mut result, right, false);
131        Ok(result)
132    }
133
134    /// Diagonal blocks of `adjoint responseᴴ`: the cotangent of each local block when
135    /// the solve's cotangent is `adjoint` and its local matrix multiplies `response`.
136    pub(crate) fn block_gradients(
137        &self,
138        adjoint: &DMatrix<Complex>,
139        response: &DMatrix<Complex>,
140    ) -> Vec<DMatrix<Complex>> {
141        let mut offset = 0;
142        self.blocks
143            .iter()
144            .map(|block| {
145                let n = block.nrows();
146                let gradient = product_views(
147                    view(adjoint).subrows(offset, n),
148                    view(response).subrows(offset, n).adjoint(),
149                );
150                offset += n;
151                gradient
152            })
153            .collect()
154    }
155}
156
157/// `I - T C` for a finite coupling of matching size.
158fn operator(local: &LocalMatrix, coupling: &DMatrix<Complex>) -> Result<DMatrix<Complex>> {
159    if coupling.shape() != (local.dimension(), local.dimension())
160        || coupling.iter().any(|z| !finite(*z))
161    {
162        return Err(Error::InvalidInput(
163            "invalid cluster coupling matrix".into(),
164        ));
165    }
166    let mut operator = -local.apply(coupling, false)?;
167    operator.set_diagonal(&(operator.diagonal().add_scalar(Complex::new(1.0, 0.0))));
168    Ok(operator)
169}
170
171/// The LU factorization of `I - T C`, kept with the local T-matrix `T` and the
172/// coupling `C`, for any number of incident-field solves.
173///
174/// A block-diagonal `T` stays a list of particle blocks; it is never expanded to a
175/// dense matrix. [`record`](Self::record) shares the factor through an [`Arc`], so
176/// every recorded solve keeps only its own incident and scattered fields.
177///
178/// treams-rs extension. treams solves the full system once with
179/// `TMatrix.interaction.solve()`; [`interaction`] computes that result.
180#[derive(Clone, Debug)]
181pub struct InteractionFactor {
182    local: LocalMatrix,
183    coupling: DMatrix<Complex>,
184    lu: Lu,
185}
186
187impl InteractionFactor {
188    /// Factor `I - T C` for a dense local matrix `T`.
189    pub fn new(local: DMatrix<Complex>, coupling: DMatrix<Complex>) -> Result<Self> {
190        Self::factorize(LocalMatrix::from_dense(local)?, coupling)
191    }
192
193    /// Factor `I - T C` for a block-diagonal `T` given by its particle blocks.
194    pub fn from_blocks(local: Vec<DMatrix<Complex>>, coupling: DMatrix<Complex>) -> Result<Self> {
195        Self::factorize(LocalMatrix::from_blocks(local)?, coupling)
196    }
197
198    fn factorize(local: LocalMatrix, coupling: DMatrix<Complex>) -> Result<Self> {
199        let lu = Lu::new(operator(&local, &coupling)?)?;
200        Ok(Self {
201            local,
202            coupling,
203            lu,
204        })
205    }
206
207    /// Number of multipole channels.
208    #[must_use]
209    pub fn dimension(&self) -> usize {
210        self.local.dimension()
211    }
212
213    /// Solve `(I - T C) scattered = T incident` for only the supplied columns.
214    pub fn solve(&self, incident: &DMatrix<Complex>) -> Result<DMatrix<Complex>> {
215        self.validate(incident)?;
216        self.solve_system(self.local.apply(incident, false)?)
217    }
218
219    /// Gradient of incident columns for a fixed factor: `Tᴴ (I - T C)⁻ᴴ G`.
220    /// This borrows the saved LU and needs no incident or scattered fields.
221    pub fn pullback_incident(&self, cotangent: &DMatrix<Complex>) -> Result<DMatrix<Complex>> {
222        self.validate(cotangent)?;
223        let mut adjoint = cotangent.clone();
224        self.lu.solve_adjoint_in_place(view_mut(&mut adjoint))?;
225        self.local.apply(&adjoint, true)
226    }
227
228    /// Solve like [`solve`](Self::solve) and keep `B + C X` for both derivative
229    /// directions. This response replaces the incident fields, and the residual
230    /// shares this factor. Use `solve` when only the scattered fields are needed.
231    pub fn record(self: &Arc<Self>, incident: DMatrix<Complex>) -> Result<IlluminateResidual> {
232        let value = self.solve(&incident)?;
233        let response = incident + product(&self.coupling, &value);
234        Ok(IlluminateResidual {
235            factor: Arc::clone(self),
236            response,
237            value,
238        })
239    }
240
241    /// `(I - T C)⁻¹ rhs`; a non-finite solution means a numerically singular operator.
242    fn solve_system(&self, mut rhs: DMatrix<Complex>) -> Result<DMatrix<Complex>> {
243        self.lu.solve_in_place(view_mut(&mut rhs))?;
244        if rhs.iter().any(|z| !finite(*z)) {
245            return Err(Error::Singular);
246        }
247        Ok(rhs)
248    }
249
250    fn validate(&self, values: &DMatrix<Complex>) -> Result<()> {
251        if values.nrows() != self.dimension()
252            || values.ncols() == 0
253            || values.iter().any(|&z| !finite(z))
254        {
255            return Err(Error::InvalidInput(
256                "require finite channel-by-illumination columns matching the factorization".into(),
257            ));
258        }
259        Ok(())
260    }
261
262    fn validate_tangents(
263        &self,
264        local: &[DMatrix<Complex>],
265        coupling: &DMatrix<Complex>,
266    ) -> Result<()> {
267        self.local.validate_tangent(local)?;
268        if coupling.shape() != self.coupling.shape() || coupling.iter().any(|&z| !finite(z)) {
269            return Err(Error::InvalidInput(
270                "coupling tangent must be finite and match the coupling matrix".into(),
271            ));
272        }
273        Ok(())
274    }
275}
276
277/// What [`interaction`] saves for its pullback: the factorization of `I - T C` and
278/// the interacting T-matrix `X`.
279#[derive(Clone, Debug)]
280pub struct InteractionResidual {
281    factor: InteractionFactor,
282    value: DMatrix<Complex>,
283}
284
285/// Solve a multiple-scattering system for the interacting T-matrix
286/// `X = (I - T C)⁻¹ T`, keeping its LU factorization for the pullback.
287///
288/// `local` is the block-diagonal T-matrix `T` of all particles, and `coupling` the
289/// singular expansion `C` between them, zero on the diagonal blocks.
290///
291/// Upstream: `treams.TMatrix.interaction.solve`.
292pub fn interaction(
293    local: DMatrix<Complex>,
294    coupling: DMatrix<Complex>,
295) -> Result<InteractionResidual> {
296    solve_interacting(LocalMatrix::from_dense(local)?, coupling)
297}
298
299/// [`interaction`] for a block-diagonal `T` given by its particle blocks, which are
300/// never expanded to a dense matrix for the factorization.
301pub(crate) fn interaction_blocks(
302    blocks: Vec<DMatrix<Complex>>,
303    coupling: DMatrix<Complex>,
304) -> Result<InteractionResidual> {
305    solve_interacting(LocalMatrix::from_blocks(blocks)?, coupling)
306}
307
308fn solve_interacting(
309    local: LocalMatrix,
310    coupling: DMatrix<Complex>,
311) -> Result<InteractionResidual> {
312    let factor = InteractionFactor::factorize(local, coupling)?;
313    let value = factor.solve_system(factor.local.to_dense()?)?;
314    Ok(InteractionResidual { factor, value })
315}
316
317/// Gradients of the inputs of [`interaction`], in the order of its arguments.
318///
319/// `L` is the type of the local-matrix gradient: one matrix from
320/// [`InteractionResidual::pullback`], or one per diagonal block from
321/// [`InteractionResidual::pullback_blocks`].
322#[derive(Clone, Debug)]
323pub struct InteractionGradient<L = DMatrix<Complex>> {
324    /// Gradient of the local T-matrix `T`.
325    pub local: L,
326    /// Gradient of the coupling `C`.
327    pub coupling: DMatrix<Complex>,
328}
329
330impl InteractionResidual {
331    pub(crate) fn local_shapes(&self) -> impl ExactSizeIterator<Item = (usize, usize)> {
332        self.factor.local.blocks.iter().map(DMatrix::shape)
333    }
334
335    /// The interacting T-matrix `X`.
336    #[must_use]
337    pub const fn value(&self) -> &DMatrix<Complex> {
338        &self.value
339    }
340
341    /// The shape of the interacting T-matrix.
342    #[must_use]
343    pub fn shape(&self) -> (usize, usize) {
344        self.value.shape()
345    }
346
347    /// Directional derivative of the interacting T-matrix with respect to dense
348    /// `T` and `C`. Reuses the forward LU for one tangent solve.
349    pub fn pushforward(
350        &self,
351        local: &DMatrix<Complex>,
352        coupling: &DMatrix<Complex>,
353    ) -> Result<DMatrix<Complex>> {
354        self.pushforward_blocks(std::slice::from_ref(local), coupling)
355    }
356
357    /// Like [`pushforward`](Self::pushforward), with one tangent per local block.
358    /// Solves `(I - T C) dX = dT (I + C X) + T dC X` without expanding `dT`.
359    pub fn pushforward_blocks(
360        &self,
361        local: &[DMatrix<Complex>],
362        coupling: &DMatrix<Complex>,
363    ) -> Result<DMatrix<Complex>> {
364        let factor = &self.factor;
365        factor.validate_tangents(local, coupling)?;
366        let mut response = product(&factor.coupling, &self.value);
367        response.set_diagonal(&(response.diagonal().add_scalar(Complex::new(1.0, 0.0))));
368        let mut rhs = factor.local.apply_tangent(local, &response)?;
369        rhs += factor.local.apply(&product(coupling, &self.value), false)?;
370        factor.solve_system(rhs)
371    }
372
373    /// Gradients of `T` and `C` from `cotangent`, the gradient of a real loss with
374    /// respect to `X`.
375    ///
376    /// The pullback solves one adjoint system with the forward LU. It has no Rayon
377    /// reduction; the LU solve and the matrix products run in faer with the worker
378    /// count of [`linalg`](crate::linalg).
379    pub fn pullback(&self, cotangent: &DMatrix<Complex>) -> Result<InteractionGradient> {
380        self.pullback_with(cotangent, |_, adjoint, response| {
381            product_adjoint_right(adjoint, response)
382        })
383    }
384
385    /// Like [`pullback`](Self::pullback), with one gradient per diagonal block of a
386    /// block-diagonal `T`, or a single one for a dense `T`.
387    pub fn pullback_blocks(
388        &self,
389        cotangent: &DMatrix<Complex>,
390    ) -> Result<InteractionGradient<Vec<DMatrix<Complex>>>> {
391        self.pullback_with(cotangent, LocalMatrix::block_gradients)
392    }
393
394    fn pullback_with<T>(
395        &self,
396        cotangent: &DMatrix<Complex>,
397        local_gradient: impl FnOnce(&LocalMatrix, &DMatrix<Complex>, &DMatrix<Complex>) -> T,
398    ) -> Result<InteractionGradient<T>> {
399        if cotangent.shape() != self.shape() || cotangent.iter().any(|z| !finite(*z)) {
400            return Err(Error::InvalidInput(
401                "invalid interacting-matrix cotangent".into(),
402            ));
403        }
404        let factor = &self.factor;
405        let mut response = product(&factor.coupling, &self.value);
406        response.set_diagonal(&(response.diagonal().add_scalar(Complex::new(1.0, 0.0))));
407        let mut adjoint = cotangent.clone();
408        factor.lu.solve_adjoint_in_place(view_mut(&mut adjoint))?;
409        let local_gradient = local_gradient(&factor.local, &adjoint, &response);
410        // Reuse the derivative workspace: response becomes T^H Y, and the
411        // adjoint becomes (T^H Y) X^H. The saved factor and primal stay immutable.
412        factor.local.apply_into(&mut response, &adjoint, true);
413        product_adjoint_right_into(&mut adjoint, &response, &self.value);
414        Ok(InteractionGradient {
415            local: local_gradient,
416            coupling: adjoint,
417        })
418    }
419}
420
421/// What [`InteractionFactor::record`] saves for its pullback: the shared factor and
422/// the total incident response `B + C X` and scattered fields of the requested columns.
423#[derive(Debug)]
424pub struct IlluminateResidual {
425    factor: Arc<InteractionFactor>,
426    response: DMatrix<Complex>,
427    value: DMatrix<Complex>,
428}
429
430/// Gradients of the local T-matrix, the coupling and the incident fields of an
431/// [`IlluminateResidual`].
432#[derive(Debug)]
433pub struct IlluminateGradient {
434    /// One gradient per local block (one for a dense local matrix).
435    pub local: Vec<DMatrix<Complex>>,
436    /// Gradient of the coupling `C`.
437    pub coupling: DMatrix<Complex>,
438    /// Gradient of the incident fields `B`.
439    pub incident: DMatrix<Complex>,
440}
441
442impl IlluminateResidual {
443    /// Shapes of the recorded local blocks, in particle order.
444    pub fn local_shapes(&self) -> impl ExactSizeIterator<Item = (usize, usize)> {
445        self.factor.local.blocks.iter().map(DMatrix::shape)
446    }
447
448    /// The scattered coefficients, one column per incident field, which the pullback
449    /// reads.
450    #[must_use]
451    pub const fn value(&self) -> &DMatrix<Complex> {
452        &self.value
453    }
454
455    /// The shape of the scattered coefficients: channels and incident fields.
456    #[must_use]
457    pub fn shape(&self) -> (usize, usize) {
458        self.value.shape()
459    }
460
461    /// Directional derivative of the requested scattered columns for tangents of
462    /// the local blocks, coupling and incident fields. Uses the shared LU once:
463    /// `(I - T C) dX = dT (B + C X) + T (dB + dC X)`.
464    pub fn pushforward(
465        &self,
466        local: &[DMatrix<Complex>],
467        coupling: &DMatrix<Complex>,
468        incident: &DMatrix<Complex>,
469    ) -> Result<DMatrix<Complex>> {
470        let factor = &self.factor;
471        factor.validate_tangents(local, coupling)?;
472        if incident.shape() != self.shape() || incident.iter().any(|&z| !finite(z)) {
473            return Err(Error::InvalidInput(
474                "incident tangent must be finite and match the incident fields".into(),
475            ));
476        }
477        let mut rhs = factor.local.apply_tangent(local, &self.response)?;
478        rhs += factor
479            .local
480            .apply(&(incident + product(coupling, &self.value)), false)?;
481        factor.solve_system(rhs)
482    }
483
484    /// Gradients of `T`, `C` and `B` from `cotangent`, the gradient of a real loss with
485    /// respect to the scattered fields.
486    ///
487    /// With `Y = (I - T C)⁻ᴴ G` from the shared LU, the gradients are `Y (B + C X)ᴴ` on
488    /// each local block, `Tᴴ Y Xᴴ` and `Tᴴ Y`. The pullback has no Rayon reduction; the
489    /// LU solve and the matrix products run in faer with the worker count of
490    /// [`linalg`](crate::linalg).
491    pub fn pullback(&self, cotangent: &DMatrix<Complex>) -> Result<IlluminateGradient> {
492        self.factor.validate(cotangent)?;
493        if cotangent.shape() != self.shape() {
494            return Err(Error::InvalidInput(
495                "cotangent shape must match scattered fields".into(),
496            ));
497        }
498        let factor = &*self.factor;
499        let mut adjoint = cotangent.clone();
500        factor.lu.solve_adjoint_in_place(view_mut(&mut adjoint))?;
501        let incident = factor.local.apply(&adjoint, true)?;
502        let local = factor.local.block_gradients(&adjoint, &self.response);
503        let coupling = product_adjoint_right(&incident, &self.value);
504        Ok(IlluminateGradient {
505            local,
506            coupling,
507            incident,
508        })
509    }
510}
511
512#[cfg(test)]
513mod tests {
514    use std::sync::Arc;
515
516    use nalgebra::DMatrix;
517    use proptest::{prelude::*, test_runner::TestCaseError};
518
519    use super::{InteractionFactor, InteractionGradient, interaction};
520    use crate::{
521        Complex,
522        test_support::{DEFAULT_CASES, complex, complex_matrix, prop_assert_close, re_dot},
523    };
524
525    /// Local blocks of sizes 1 to 3, a coupling, incident fields and cotangents for
526    /// one to four illuminations.
527    type Inputs = (
528        Vec<DMatrix<Complex>>,
529        DMatrix<Complex>,
530        DMatrix<Complex>,
531        DMatrix<Complex>,
532    );
533
534    fn inputs() -> impl Strategy<Value = Inputs> {
535        (prop::collection::vec(1_usize..=3, 1..=3), 1_usize..=4).prop_flat_map(|(sizes, p)| {
536            let n = sizes.iter().sum();
537            (
538                sizes
539                    .iter()
540                    .map(|&size| complex_matrix(size, size, 1.0))
541                    .collect::<Vec<_>>(),
542                complex_matrix(n, n, 1.0),
543                complex_matrix(n, p, 1.0),
544                complex_matrix(n, p, 1.0),
545            )
546        })
547    }
548
549    proptest! {
550        // Each case factors three systems of at most 9 channels.
551        #![proptest_config(ProptestConfig::with_cases(DEFAULT_CASES))]
552
553        #[test]
554        fn requested_columns_match_the_implicit_adjoint_closed_forms(
555            (blocks, coupling, incident, weight) in inputs(),
556            amplitude in complex(2.0),
557        ) {
558            check_requested_columns(blocks, &coupling, &incident, &weight, amplitude)?;
559        }
560    }
561
562    /// With `X = (I - T C)⁻¹ T A` and `Y = (I - T C)⁻ᴴ G` from an independent nalgebra
563    /// LU, recorded solves equal `X` and are linear in `A`, and the pullback returns
564    /// the closed forms of the implicit adjoint (formal/Formal/ImplicitAdjoint.lean):
565    /// `T̄ = Y (A + C X)ᴴ`, restricted to the local blocks, `C̄ = Tᴴ Y Xᴴ` and
566    /// `Ā = Tᴴ Y`. The dense factorization returns the full `T̄`, and the interacting
567    /// T-matrix applied to `A`, with the cotangent `G Aᴴ`, gives the same gradients.
568    fn check_requested_columns(
569        mut blocks: Vec<DMatrix<Complex>>,
570        coupling: &DMatrix<Complex>,
571        incident: &DMatrix<Complex>,
572        weight: &DMatrix<Complex>,
573        amplitude: Complex,
574    ) -> Result<(), TestCaseError> {
575        let n = coupling.nrows();
576        let mut local = DMatrix::zeros(n, n);
577        let mut offset = 0;
578        for block in &blocks {
579            local
580                .view_mut((offset, offset), block.shape())
581                .copy_from(block);
582            offset += block.nrows();
583        }
584        // ‖T‖ ‖C‖ < 1/2 keeps I - T C well conditioned: ‖(I - T C)⁻¹‖ < 2.
585        let shrink = Complex::from(0.5 / (1.0 + local.norm() * coupling.norm()));
586        local *= shrink;
587        for block in &mut blocks {
588            *block *= shrink;
589        }
590        let identity = DMatrix::identity(n, n);
591        let system = &identity - &local * coupling;
592        let x = system.clone().lu().solve(&(&local * incident)).unwrap();
593        let y = system.adjoint().lu().solve(weight).unwrap();
594        let full_local = &y * (incident + coupling * &x).adjoint();
595        let expected_coupling = local.adjoint() * &y * x.adjoint();
596        let expected_incident = local.adjoint() * &y;
597        let value_tolerance = 1e-12 * (1.0 + incident.norm());
598        let tolerance = value_tolerance * (1.0 + weight.norm());
599
600        let factor =
601            Arc::new(InteractionFactor::from_blocks(blocks.clone(), coupling.clone()).unwrap());
602        let residual = factor.record(incident.clone()).unwrap();
603        prop_assert_close!(residual.value(), &x, value_tolerance);
604        let scaled = factor.solve(&(incident * amplitude)).unwrap();
605        let expected = residual.value() * amplitude;
606        prop_assert_close!(scaled, expected, value_tolerance * (1.0 + amplitude.norm()));
607        let gradient = residual.pullback(weight).unwrap();
608        prop_assert_eq!(gradient.local.len(), blocks.len());
609        let mut offset = 0;
610        for (block, actual) in blocks.iter().zip(&gradient.local) {
611            let expected = full_local
612                .view((offset, offset), block.shape())
613                .into_owned();
614            prop_assert_close!(actual, &expected, tolerance, "block at {}", offset);
615            offset += block.nrows();
616        }
617        prop_assert_close!(&gradient.coupling, &expected_coupling, tolerance);
618        prop_assert_close!(&gradient.incident, &expected_incident, tolerance);
619        let incident_gradient = factor.pullback_incident(weight).unwrap();
620        prop_assert_close!(&incident_gradient, &expected_incident, tolerance);
621        prop_assert_close!(
622            re_dot(weight, residual.value()),
623            re_dot(&incident_gradient, incident),
624            tolerance * (1.0 + incident.norm())
625        );
626
627        let dense = Arc::new(InteractionFactor::new(local.clone(), coupling.clone()).unwrap());
628        let dense_gradient = dense
629            .record(incident.clone())
630            .unwrap()
631            .pullback(weight)
632            .unwrap();
633        prop_assert_eq!(dense_gradient.local.len(), 1);
634        prop_assert_close!(&dense_gradient.local[0], &full_local, tolerance);
635        prop_assert_close!(
636            &dense.pullback_incident(weight).unwrap(),
637            &incident_gradient,
638            tolerance
639        );
640
641        let interacting = interaction(local, coupling.clone()).unwrap();
642        prop_assert_close!(interacting.value() * incident, x, value_tolerance);
643        let InteractionGradient {
644            local: local_gradient,
645            coupling: coupling_gradient,
646        } = interacting
647            .pullback(&(weight * incident.adjoint()))
648            .unwrap();
649        prop_assert_close!(local_gradient, full_local, tolerance);
650        prop_assert_close!(coupling_gradient, expected_coupling, tolerance);
651        Ok(())
652    }
653}
treams_core/cluster/interaction.rs

treams_core/cluster/
interaction.rs
