kryst 3.2.1

Krylov subspace and preconditioned iterative solvers for dense and sparse linear systems, with shared and distributed memory parallelism.
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; // silence unused warning
            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)
    }
}