kryst/solver/
mod.rs

1//! Solver traits and implementations.
2//!
3//! Most users configure solvers through [`KspContext`](crate::context::KspContext) and
4//! [`SolverType`](crate::context::ksp_context::SolverType). This module exposes the
5//! underlying solver types and traits for advanced use, custom pipelines, or testing.
6
7use crate::context::ksp_context::Workspace;
8use crate::error::KError;
9use crate::matrix::op::LinOp;
10use crate::parallel::UniverseComm;
11use crate::preconditioner::{PcSide, Preconditioner};
12use crate::utils::convergence::SolveStats;
13use std::any::Any;
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub enum MonitorAction {
17    Continue,
18    Stop,
19}
20
21pub type MonitorCallback<R> = dyn Fn(usize, R, usize) -> MonitorAction + Send + Sync;
22
23pub mod api;
24pub use api::Solver;
25pub mod adapters;
26pub use crate::ops::klinop::KLinOp;
27pub use crate::ops::kpc::KPreconditioner;
28pub use adapters::LegacyDirectAdapter;
29
30pub mod block;
31
32/// Object-safe linear solver operating on `f64` slices and [`LinOp`] operators.
33pub trait LinearSolver: Send + Any {
34    type Error;
35
36    fn as_any_mut(&mut self) -> &mut dyn Any;
37
38    /// Allow solver to configure workspace buffers.
39    fn setup_workspace(&mut self, _work: &mut Workspace) {}
40
41    /// Solve `a * x = b` optionally using a preconditioner.
42    fn solve(
43        &mut self,
44        a: &dyn LinOp<S = f64>,
45        pc: Option<&mut dyn Preconditioner>,
46        b: &[f64],
47        x: &mut [f64],
48        pc_side: PcSide,
49        comm: &UniverseComm,
50        monitors: Option<&[Box<MonitorCallback<f64>>]>,
51        work: Option<&mut Workspace>,
52    ) -> Result<SolveStats<f64>, Self::Error>;
53}
54
55/// Legacy generic solver trait retained for existing implementations.
56pub mod legacy {
57    use crate::algebra::prelude::KrystScalar;
58    use crate::preconditioner::legacy::Preconditioner;
59    use crate::solver::MonitorCallback;
60    use crate::utils::convergence::SolveStats;
61
62    pub trait LinearSolver<M: ?Sized, V> {
63        type Error;
64        type Scalar: KrystScalar;
65
66        fn solve(
67            &mut self,
68            a: &M,
69            pc: Option<&(dyn Preconditioner<M, V> + '_)>,
70            b: &V,
71            x: &mut V,
72            pc_side: crate::preconditioner::PcSide,
73            comm: &crate::parallel::UniverseComm,
74            monitors: Option<&[Box<MonitorCallback<Self::Scalar>>]>,
75            work: Option<&mut crate::context::ksp_context::Workspace>,
76        ) -> Result<SolveStats<Self::Scalar>, Self::Error>;
77
78        fn setup_workspace(&mut self, _work: &mut crate::context::ksp_context::Workspace) {}
79
80        fn solve_simple(
81            &mut self,
82            a: &M,
83            pc: Option<&(dyn Preconditioner<M, V> + '_)>,
84            b: &V,
85            x: &mut V,
86            pc_side: crate::preconditioner::PcSide,
87            comm: &crate::parallel::UniverseComm,
88        ) -> Result<SolveStats<Self::Scalar>, Self::Error>
89        where
90            Self: Sized,
91        {
92            self.solve(a, pc, b, x, pc_side, comm, None, None)
93        }
94
95        fn solve_with_monitors(
96            &mut self,
97            a: &M,
98            pc: Option<&(dyn Preconditioner<M, V> + '_)>,
99            b: &V,
100            x: &mut V,
101            pc_side: crate::preconditioner::PcSide,
102            comm: &crate::parallel::UniverseComm,
103            monitors: &[Box<MonitorCallback<Self::Scalar>>],
104        ) -> Result<SolveStats<Self::Scalar>, Self::Error>
105        where
106            Self: Sized,
107        {
108            self.solve(a, pc, b, x, pc_side, comm, Some(monitors), None)
109        }
110
111        fn solve_with_workspace(
112            &mut self,
113            a: &M,
114            pc: Option<&(dyn Preconditioner<M, V> + '_)>,
115            b: &V,
116            x: &mut V,
117            pc_side: crate::preconditioner::PcSide,
118            comm: &crate::parallel::UniverseComm,
119            work: &mut crate::context::ksp_context::Workspace,
120        ) -> Result<SolveStats<Self::Scalar>, Self::Error>
121        where
122            Self: Sized,
123        {
124            self.solve(a, pc, b, x, pc_side, comm, None, Some(work))
125        }
126    }
127}
128
129/// Adapter allowing legacy generic solvers (MatVec/Vec) to be used with the
130/// new object-safe [`LinearSolver`] trait over [`LinOp`].
131pub struct OpSolverAdapter<S> {
132    inner: S,
133}
134
135impl<S> OpSolverAdapter<S> {
136    pub fn new(inner: S) -> Self {
137        Self { inner }
138    }
139
140    pub fn inner_mut(&mut self) -> &mut S {
141        &mut self.inner
142    }
143}
144
145#[cfg(not(feature = "complex"))]
146struct OpPcAdapter<'p> {
147    inner: &'p dyn Preconditioner,
148}
149
150#[cfg(not(feature = "complex"))]
151impl<'p, 'm> crate::preconditioner::legacy::Preconditioner<dyn LinOp<S = f64> + 'm, Vec<f64>>
152    for OpPcAdapter<'p>
153{
154    fn setup(&mut self, _a: &(dyn LinOp<S = f64> + 'm)) -> Result<(), KError> {
155        Ok(())
156    }
157    fn apply(&self, side: PcSide, r: &Vec<f64>, z: &mut Vec<f64>) -> Result<(), KError> {
158        self.inner.apply(side, r.as_slice(), z.as_mut_slice())
159    }
160}
161
162#[cfg(not(feature = "complex"))]
163impl<S> LinearSolver for OpSolverAdapter<S>
164where
165    S: for<'a> legacy::LinearSolver<
166            dyn LinOp<S = f64> + 'a,
167            Vec<f64>,
168            Scalar = f64,
169            Error = KError,
170        > + Send
171        + 'static,
172{
173    type Error = KError;
174
175    fn as_any_mut(&mut self) -> &mut dyn Any {
176        self
177    }
178
179    fn setup_workspace(&mut self, work: &mut Workspace) {
180        self.inner.setup_workspace(work);
181    }
182
183    fn solve(
184        &mut self,
185        a: &dyn LinOp<S = f64>,
186        pc: Option<&mut dyn Preconditioner>,
187        b: &[f64],
188        x: &mut [f64],
189        pc_side: PcSide,
190        comm: &UniverseComm,
191        monitors: Option<&[Box<MonitorCallback<f64>>]>,
192        work: Option<&mut Workspace>,
193    ) -> Result<SolveStats<f64>, Self::Error> {
194        let mut x_vec = x.to_vec();
195        let b_vec = b.to_vec();
196        // Coerce &mut dyn Preconditioner -> &dyn Preconditioner for the legacy adapter.
197        let pc_adapter = pc.as_deref().map(|p| OpPcAdapter { inner: p });
198        let pc_ref = pc_adapter.as_ref().map(|p| {
199            p
200                as &dyn crate::preconditioner::legacy::Preconditioner<
201                    dyn LinOp<S = f64> + '_,
202                    Vec<f64>,
203                >
204        });
205        let stats = self
206            .inner
207            .solve(a, pc_ref, &b_vec, &mut x_vec, pc_side, comm, monitors, work)?;
208        x.copy_from_slice(&x_vec);
209        Ok(stats)
210    }
211}
212
213// Re-export solver implementations
214pub mod cg;
215pub use cg::CgSolver;
216pub mod cgnr;
217pub use cgnr::CgnrSolver;
218
219pub mod gmres;
220pub use gmres::GmresSolver;
221pub mod fgmres;
222pub use fgmres::FgmresSolver;
223pub mod bicgstab;
224pub use bicgstab::{BiCgStabBreakdownPolicy, BiCgStabSolver, BiCgStabVariant};
225pub mod idrs;
226pub use idrs::{
227    BreakdownRepair as IdrsBreakdownRepair, IdrsBuilder, IdrsOptions, IdrsSolver,
228    Omega as IdrsOmega, ShadowP as IdrsShadowP,
229};
230pub mod cgs;
231pub use cgs::CgsSolver;
232pub mod richardson;
233pub use richardson::RichardsonSolver;
234pub mod chebyshev;
235pub use chebyshev::ChebyshevSolver;
236pub mod cr;
237pub use cr::CrSolver;
238pub mod pcg;
239pub use pcg::{PCG_PIPELINED_DEFAULT_REPLACE_EVERY, PcgSolver, PcgVariant};
240pub mod minres;
241pub use minres::MinresSolver;
242pub mod lsmr;
243pub use lsmr::LsmrSolver;
244pub mod lsqr;
245pub use lsqr::LsqrSolver;
246// Dense direct modules are gated; opt-in via `dense-direct`.
247#[cfg(feature = "dense-direct")]
248pub mod dense_lu;
249#[cfg(feature = "dense-direct")]
250pub mod dense_qr;
251#[cfg(feature = "backend-faer")]
252pub mod direct_lu;
253
254#[cfg(feature = "backend-faer")]
255pub use direct_lu::{LuSolver, QrSolver};
256#[cfg(feature = "superlu_dist")]
257pub mod superlu_dist;
258#[cfg(feature = "superlu_dist")]
259pub use superlu_dist::SuperLuDistSolver;
260
261pub mod qmr;
262pub use qmr::QmrSolver;
263pub mod tfqmr;
264pub use tfqmr::TfqmrSolver;
265pub mod tcqmr;
266pub use tcqmr::TcqmrSolver;
267
268pub mod gcr;
269pub use gcr::GcrSolver;
270pub mod pipegcr;
271pub use pipegcr::PipeGcrSolver;
272
273pub mod pca_gmres;
274pub use pca_gmres::{PcaGmresSolver, PcaPcMode};
275
276pub mod common;
277
278#[cfg(test)]
279mod tests;
kryst/solver/mod.rs

kryst/solver/
mod.rs
