Source code for ossdbs.fem.preconditioner

# Copyright 2023, 2024 Johannes Reding, Julius Zimmermann
# SPDX-License-Identifier: GPL-3.0-or-later

import logging
from abc import ABC, abstractmethod

import ngsolve
import numpy as np

_logger = logging.getLogger(__name__)


class Preconditioner(ABC):
    """Define a preconditioner."""

    @abstractmethod
    def to_dictionary(self) -> dict:
        """Prepare dictionary to be passed to NGSolve."""
        pass


class BDDCPreconditioner(Preconditioner):
    """BDDC preconditioner, highly efficient for higher-order FEM."""

    def __init__(self, coarsetype: str = "h1amg") -> None:
        _logger.debug(f"Initialized BDDC preconditioner with coarsetype: {coarsetype}")
        self.type = "bddc"
        self.coarsetype = coarsetype

    def to_dictionary(self) -> dict:
        """Prepare dictionary to be passed to NGSolve."""
        return {"type": self.type, "coarsetype": self.coarsetype}


class LocalPreconditioner(Preconditioner):
    """Jacobi preconditioner."""

    def __init__(self) -> None:
        _logger.debug("Initialized Jacobi preconditioner")
        self.type = "local"

    def to_dictionary(self) -> dict:
        """Prepare dictionary to be passed to NGSolve."""
        return {"type": self.type}


class CustomizedLocalPreconditioner(Preconditioner):
    """Custom diagonal Jacobi preconditioner."""

    def __init__(self) -> None:
        _logger.debug("Initialized customized Jacobi preconditioner")
        self.type = "customized_local"

    def to_dictionary(self) -> dict:
        """Prepare dictionary to be passed to the solver."""
        return {"type": self.type}


class JacobiPreconditioner(ngsolve.BaseMatrix):
    """Simple diagonal Jacobi preconditioner as an NGSolve operator."""

    def __init__(self, inv_diag: np.ndarray, freedofs, template_vec) -> None:
        super().__init__()
        self._inv_diag = inv_diag
        self._freedofs = np.array(freedofs, dtype=bool)
        self._template_vec = template_vec

    def Mult(self, x, y) -> None:
        """Apply the preconditioner to ``x`` and store the result in ``y``."""
        x_np = x.FV().NumPy()
        y_np = y.FV().NumPy()
        y_np[:] = self._inv_diag * x_np
        y_np[~self._freedofs] = 0.0

    def CreateColVector(self):
        """Create a compatible column vector."""
        return self._template_vec.CreateVector()

    def CreateRowVector(self):
        """Create a compatible row vector."""
        return self._template_vec.CreateVector()

    @property
    def height(self) -> int:
        """Return the operator height."""
        return len(self._inv_diag)

    @property
    def width(self) -> int:
        """Return the operator width."""
        return len(self._inv_diag)


class MultigridPreconditioner(Preconditioner):
    """Geometric multigrid preconditioner."""

    def __init__(self) -> None:
        _logger.debug("Initialized multigrid preconditioner")
        self.type = "multigrid"

    def to_dictionary(self) -> dict:
        """Prepare dictionary to be passed to NGSolve."""
        return {"type": self.type}


class AMGPreconditioner(Preconditioner):
    """Algebraic multigrid preconditioner."""

    def __init__(self) -> None:
        _logger.debug("Initialized AMG preconditioner")
        self.type = "h1amg"

    def to_dictionary(self) -> dict:
        """Prepare dictionary to be passed to NGSolve."""
        return {"type": self.type}


class DirectPreconditioner(Preconditioner):
    """Direct solver."""

    def __init__(self, inverse: str = "") -> None:
        _logger.debug("Use direct solver as preconditioner")
        self.type = "direct"
        self.inverse = inverse

    def to_dictionary(self) -> dict:
        """Prepare dictionary to be passed to NGSolve."""
        if self.inverse == "":
            return {"type": self.type}
        return {"type": self.type, "inverse": self.inverse}