Skip to main content

rlst/operator/interface/
matrix_operator.rs

1//! Matrix interface.
2//!
3//! This module defines interfaces to the Rlst dense and sparse matrix types
4//! to be used as abstract operators.
5
6#[cfg(feature = "mpi")]
7use mpi::traits::{Communicator, Equivalence};
8
9use crate::{
10    Array, AsMatrixApply, BaseItem, LinearSpace, Shape, abstract_operator::OperatorBase,
11    dense::array::DynArray, operator::abstract_operator::Operator, sparse::csr_mat::CsrMatrix,
12};
13
14#[cfg(feature = "mpi")]
15use crate::{
16    dense::{base_array::BaseArray, data_container::VectorContainer},
17    operator::interface::distributed_array_vector_space::DistributedArrayVectorSpace,
18    sparse::distributed_array::DistributedArray,
19    sparse::distributed_csr_mat::DistributedCsrMatrix,
20};
21
22use super::array_vector_space::ArrayVectorSpace;
23
24/// Borrowed dense matrix viewed as an abstract operator.
25///
26/// For an `m x n` matrix, the domain contains length-n vectors and the range
27/// length-m vectors. The operator cannot outlive the matrix.
28pub struct ArrayOperator<'a, ArrayImpl>
29where
30    ArrayImpl: BaseItem,
31{
32    op: &'a Array<ArrayImpl, 2>,
33    domain: ArrayVectorSpace<ArrayImpl::Item>,
34    range: ArrayVectorSpace<ArrayImpl::Item>,
35}
36
37impl<'a, ArrayImpl: BaseItem> ArrayOperator<'a, ArrayImpl> {
38    /// Borrow a matrix and construct matching domain and range spaces.
39    pub fn new(op: &'a Array<ArrayImpl, 2>) -> Self
40    where
41        ArrayImpl: Shape<2>,
42    {
43        let shape = op.shape();
44        Self {
45            op,
46            domain: ArrayVectorSpace::new(shape[1]),
47            range: ArrayVectorSpace::new(shape[0]),
48        }
49    }
50}
51
52impl<'a, ArrayImpl: BaseItem + Shape<2>> From<&'a Array<ArrayImpl, 2>>
53    for Operator<ArrayOperator<'a, ArrayImpl>>
54where
55    ArrayOperator<'a, ArrayImpl>: OperatorBase,
56{
57    fn from(value: &'a Array<ArrayImpl, 2>) -> Self {
58        Operator::new(ArrayOperator::new(value))
59    }
60}
61
62impl<'a, ArrayImpl> OperatorBase for ArrayOperator<'a, ArrayImpl>
63where
64    ArrayImpl: BaseItem,
65    ArrayVectorSpace<ArrayImpl::Item>:
66        LinearSpace<F = ArrayImpl::Item, Impl = DynArray<ArrayImpl::Item, 1>>,
67    Array<ArrayImpl, 2>: AsMatrixApply<DynArray<ArrayImpl::Item, 1>, DynArray<ArrayImpl::Item, 1>>
68        + BaseItem<Item = ArrayImpl::Item>,
69{
70    type Domain = ArrayVectorSpace<ArrayImpl::Item>;
71
72    type Range = ArrayVectorSpace<ArrayImpl::Item>;
73
74    fn domain(&self) -> &Self::Domain {
75        &self.domain
76    }
77
78    fn range(&self) -> &Self::Range {
79        &self.range
80    }
81
82    #[inline(always)]
83    fn apply(
84        &self,
85        alpha: <Self::Range as crate::LinearSpace>::F,
86        x: &crate::operator::element::Element<Self::Domain>,
87        beta: <Self::Range as crate::LinearSpace>::F,
88        y: &mut crate::operator::element::Element<Self::Range>,
89    ) {
90        self.op.apply(alpha, x.imp(), beta, y.imp_mut());
91    }
92}
93
94/// Borrowed CSR matrix viewed as an abstract operator.
95///
96/// For an `m x n` matrix, the domain contains length-n vectors and the range
97/// length-m vectors.
98pub struct CsrOperator<'a, Item> {
99    arr: &'a CsrMatrix<Item>,
100    domain: ArrayVectorSpace<Item>,
101    range: ArrayVectorSpace<Item>,
102}
103
104impl<'a, Item> CsrOperator<'a, Item> {
105    /// Borrow a CSR matrix and construct matching domain and range spaces.
106    pub fn new(arr: &'a CsrMatrix<Item>) -> Self {
107        let shape = arr.shape();
108        Self {
109            arr,
110            domain: ArrayVectorSpace::new(shape[1]),
111            range: ArrayVectorSpace::new(shape[0]),
112        }
113    }
114}
115
116impl<'a, Item> From<&'a CsrMatrix<Item>> for Operator<CsrOperator<'a, Item>>
117where
118    CsrOperator<'a, Item>: OperatorBase,
119{
120    fn from(value: &'a CsrMatrix<Item>) -> Self {
121        Operator::new(CsrOperator::new(value))
122    }
123}
124
125impl<'a, Item> OperatorBase for CsrOperator<'a, Item>
126where
127    ArrayVectorSpace<Item>: LinearSpace<F = Item, Impl = DynArray<Item, 1>>,
128    CsrMatrix<Item>: AsMatrixApply<DynArray<Item, 1>, DynArray<Item, 1>> + BaseItem<Item = Item>,
129{
130    type Domain = ArrayVectorSpace<Item>;
131
132    type Range = ArrayVectorSpace<Item>;
133
134    fn domain(&self) -> &Self::Domain {
135        &self.domain
136    }
137
138    fn range(&self) -> &Self::Range {
139        &self.range
140    }
141
142    #[inline(always)]
143    fn apply(
144        &self,
145        alpha: <Self::Range as LinearSpace>::F,
146        x: &crate::operator::element::Element<Self::Domain>,
147        beta: <Self::Range as LinearSpace>::F,
148        y: &mut crate::operator::element::Element<Self::Range>,
149    ) {
150        self.arr.apply(alpha, x.imp(), beta, y.imp_mut());
151    }
152}
153
154#[cfg(feature = "mpi")]
155/// Borrowed distributed CSR matrix viewed as an abstract operator.
156///
157/// Domain and range reuse the matrix's distributed index layouts; application
158/// is collective whenever the underlying distributed matrix operation is.
159pub struct DistributedCsrOperator<'a, Item: Equivalence, C: Communicator> {
160    op: &'a DistributedCsrMatrix<'a, Item, C>,
161    domain: DistributedArrayVectorSpace<'a, C, Item>,
162    range: DistributedArrayVectorSpace<'a, C, Item>,
163}
164
165#[cfg(feature = "mpi")]
166impl<'a, Item: Copy + Equivalence, C: Communicator> DistributedCsrOperator<'a, Item, C> {
167    /// Borrow a distributed matrix and clone its domain/range layout handles.
168    pub fn new(op: &'a DistributedCsrMatrix<'a, Item, C>) -> Self {
169        Self {
170            op,
171            domain: DistributedArrayVectorSpace::new(op.domain_layout().clone()),
172            range: DistributedArrayVectorSpace::new(op.range_layout().clone()),
173        }
174    }
175}
176
177#[cfg(feature = "mpi")]
178impl<'a, Item: Copy + Equivalence, C: Communicator> From<&'a DistributedCsrMatrix<'a, Item, C>>
179    for Operator<DistributedCsrOperator<'a, Item, C>>
180where
181    DistributedCsrOperator<'a, Item, C>: OperatorBase,
182{
183    fn from(value: &'a DistributedCsrMatrix<'a, Item, C>) -> Self {
184        Operator::new(DistributedCsrOperator::new(value))
185    }
186}
187
188#[cfg(feature = "mpi")]
189impl<'a, Item: Equivalence, C: Communicator> OperatorBase for DistributedCsrOperator<'a, Item, C>
190where
191    DistributedArrayVectorSpace<'a, C, Item>: LinearSpace<
192            F = Item,
193            Impl = DistributedArray<'a, C, BaseArray<VectorContainer<Item>, 1>, 1>,
194        >,
195    DistributedCsrMatrix<'a, Item, C>: AsMatrixApply<
196            DistributedArray<'a, C, BaseArray<VectorContainer<Item>, 1>, 1>,
197            DistributedArray<'a, C, BaseArray<VectorContainer<Item>, 1>, 1>,
198        > + BaseItem<Item = Item>,
199{
200    type Domain = DistributedArrayVectorSpace<'a, C, Item>;
201
202    type Range = DistributedArrayVectorSpace<'a, C, Item>;
203
204    fn domain(&self) -> &Self::Domain {
205        &self.domain
206    }
207
208    fn range(&self) -> &Self::Range {
209        &self.range
210    }
211
212    #[inline(always)]
213    fn apply(
214        &self,
215        alpha: <Self::Range as LinearSpace>::F,
216        x: &crate::operator::element::Element<Self::Domain>,
217        beta: <Self::Range as LinearSpace>::F,
218        y: &mut crate::operator::element::Element<Self::Range>,
219    ) {
220        self.op.apply(alpha, x.imp(), beta, y.imp_mut());
221    }
222}