1#[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
24pub 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 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
94pub 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 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")]
155pub 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 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}