use crate::algebra::prelude::*;
use crate::error::KError;
#[cfg(feature = "superlu_dist")]
use crate::matrix::convert::csr_from_linop;
use crate::matrix::op::LinOp;
use crate::parallel::UniverseComm;
use crate::preconditioner::{PcSide, Preconditioner};
#[cfg(feature = "complex")]
use crate::algebra::bridge::BridgeScratch;
#[cfg(feature = "complex")]
use crate::ops::kpc::KPreconditioner;
#[cfg(feature = "complex")]
use crate::preconditioner::bridge::apply_pc_s;
#[cfg_attr(docsrs, doc(cfg(feature = "superlu_dist")))]
pub struct SuperLuDistPc {
#[allow(dead_code)]
comm: Option<UniverseComm>,
}
impl Default for SuperLuDistPc {
fn default() -> Self {
Self::new()
}
}
impl SuperLuDistPc {
pub fn new() -> Self {
Self { comm: None }
}
}
impl Preconditioner for SuperLuDistPc {
fn setup(&mut self, pmat: &dyn LinOp<S = S>) -> Result<(), KError> {
#[cfg(feature = "superlu_dist")]
{
self.comm = Some(pmat.comm());
let _ = csr_from_linop(pmat, 0.0)?;
Ok(())
}
#[cfg(not(feature = "superlu_dist"))]
{
let _ = pmat; Err(KError::SolveError(
"superlu_dist feature not enabled".into(),
))
}
}
fn apply(&self, _side: PcSide, r: &[S], z: &mut [S]) -> Result<(), KError> {
let _ = (r, z);
Err(KError::Unsupported(
"SuperLuDistPc::apply is PREONLY-only; use SolverType::Preonly or call direct_solve",
))
}
fn direct_solve(
&mut self,
pmat: &dyn LinOp<S = S>,
b: &[S],
x: &mut [S],
) -> Result<(), KError> {
#[cfg(feature = "superlu_dist")]
{
let a = csr_from_linop(pmat, 0.0)?;
let comm = self.comm.clone().unwrap_or_else(|| pmat.comm());
#[cfg(not(feature = "complex"))]
{
crate::solver::superlu_dist::solve(a.as_ref(), b, x, &comm)
}
#[cfg(feature = "complex")]
{
let n = b.len();
let mut b_re = vec![0.0; n];
let mut b_im = vec![0.0; n];
for (i, bi) in b.iter().enumerate() {
b_re[i] = bi.real();
b_im[i] = bi.imag();
}
let mut x_re = vec![0.0; x.len()];
let mut x_im = vec![0.0; x.len()];
crate::solver::superlu_dist::solve(a.as_ref(), &b_re, &mut x_re, &comm)?;
crate::solver::superlu_dist::solve(a.as_ref(), &b_im, &mut x_im, &comm)?;
for (xi, (&re, &im)) in x.iter_mut().zip(x_re.iter().zip(x_im.iter())) {
*xi = S::from_parts(re, im);
}
Ok(())
}
}
#[cfg(not(feature = "superlu_dist"))]
{
let _ = (pmat, b, x);
Err(KError::SolveError(
"superlu_dist feature not enabled".into(),
))
}
}
fn required_format(&self) -> crate::matrix::format::OpFormat {
crate::matrix::format::OpFormat::Csr
}
}
#[cfg(feature = "complex")]
impl KPreconditioner for SuperLuDistPc {
type Scalar = S;
#[inline]
fn dims(&self) -> (usize, usize) {
<Self as Preconditioner>::dims(self)
}
#[inline]
fn apply_s(
&self,
side: PcSide,
x: &[S],
y: &mut [S],
scratch: &mut BridgeScratch,
) -> Result<(), KError> {
apply_pc_s(self, side, x, y, scratch)
}
}