<commit_before>import warnings
import scipy.sparse

try:
    import pyamg
    has_pyamg = True
except ImportError:
    warnings.warn("pyamg was not found, falling back to SciPy CG solver")
    import scipy.sparse.linalg as spalg
    has_pyamg = False

from .. import cyclic
from ... import veros_method


@veros_method
def solve(veros, rhs, sol, boundary_val=None):
    """
    Main solver for streamfunction. Solves a 2D Poisson equation. Uses either pyamg
    or scipy.sparse.linalg linear solvers.

    Arguments:
        rhs: Right-hand side vector
        sol: Initial guess, gets overwritten with solution
        boundary_val: Array containing values to set on boundary elements. Defaults to `sol`.
    """
    veros.flush()
    # only initialize solver if parent object changes
    if not solve.veros or solve.veros != id(veros):
        if has_pyamg:
            solve.linear_solver = _get_amg_solver(veros)
        else:
            solve.linear_solver = _get_scipy_solver(veros)
        solve.veros = id(veros)

    if veros.enable_cyclic_x:
        cyclic.setcyclic_x(sol)

    if boundary_val is None:
        boundary_val = sol

    z = np.prod(~veros.boundary_mask, axis=2).astype(np.bool)
    rhs[...] = np.where(z, rhs, boundary_val) # set right hand side on boundaries
    linear_solution = solve.linear_solver(rhs, sol)
    sol[...] = boundary_val
    sol[2:-2, 2:-2] = linear_solution.reshape(veros.nx + 4, veros.ny + 4)[2:-2, 2:-2]


solve.veros = None


@veros_method
def _get_scipy_solver(veros):
    matrix = _assemble_poisson_matrix(veros)
    preconditioner = _jacobi_preconditioner(veros, matrix)
    preconditioner.diagonal()[...] *= np.prod(~veros.boundary_mask, axis=2).astype(np.int).flatten()
    matrix = preconditioner * matrix

    def scipy_solver(rhs, x0):
        rhs = rhs.flatten() * preconditioner.diagonal()
        solution, info = spalg.bicgstab(matrix, rhs,
                                        x0=x0.flatten(), tol=veros.congr_epsilon,
                                        maxiter=veros.congr_max_iterations)
        if info > 0:
            warnings.warn("Streamfunction solver did not converge after {} iterations".format(info))
        return solution
    return scipy_solver


@veros_method
def _get_amg_solver(veros):
    matrix = _assemble_poisson_matrix(veros)
    if veros.backend_name == "bohrium":
        near_null_space = np.ones(matrix.shape[0], bohrium=False, dtype=veros.default_float_type)
    else:
        near_null_space = np.ones(matrix.shape[0], dtype=veros.default_float_type)
    ml = pyamg.smoothed_aggregation_solver(matrix, near_null_space[:, np.newaxis])

    def amg_solver(rhs, x0):
        if veros.backend_name == "bohrium":
            rhs = rhs.copy2numpy()
            x0 = x0.copy2numpy()
        residuals = []
        tolerance = veros.congr_epsilon * 1e-8  # to achieve the same precision as the preconditioned scipy solver
        solution = ml.solve(b=rhs.flatten(), x0=x0.flatten(), tol=tolerance,
                            residuals=residuals, accel="bicgstab")
        rel_res = residuals[-1] / residuals[0]
        if rel_res > tolerance:
            warnings.warn(
                "Streamfunction solver did not converge - residual: {:.2e}".format(rel_res))
        return np.asarray(solution)
    return amg_solver


@veros_method
def _jacobi_preconditioner(veros, matrix):
    """
    Construct a simple Jacobi preconditioner
    """
    Z = np.ones((veros.nx + 4, veros.ny + 4), dtype=veros.default_float_type)
    Y = matrix.diagonal().copy().reshape(veros.nx + 4, veros.ny + 4)[2:-2, 2:-2]
    Z[2:-2, 2:-2] = np.where(Y != 0., 1. / Y, 1.)
    return scipy.sparse.dia_matrix((Z.flatten(), 0), shape=(Z.size, Z.size)).tocsr()


@veros_method
def _assemble_poisson_matrix(veros):
    """
    Construct a sparse matrix based on the stencil for the 2D Poisson equation.
    """
    # assemble diagonals
    main_diag = np.ones((veros.nx + 4, veros.ny + 4), dtype=veros.default_float_type)<commit_msg>[core] work on using pyamg as preconditioner<commit_after>import warnings
import scipy.sparse

import scipy.sparse.linalg as spalg
try:
    import pyamg
    has_pyamg = True
except ImportError:
    warnings.warn("pyamg was not found, falling back to SciPy CG solver")
    has_pyamg = False

from .. import cyclic
from ... import veros_method


@veros_method
def solve(veros, rhs, sol, boundary_val=None):
    """
    Main solver for streamfunction. Solves a 2D Poisson equation. Uses either pyamg
    or scipy.sparse.linalg linear solvers.

    Arguments:
        rhs: Right-hand side vector
        sol: Initial guess, gets overwritten with solution
        boundary_val: Array containing values to set on boundary elements. Defaults to `sol`.
    """
    veros.flush()
    # only initialize solver if parent object changes
    if not solve.veros or solve.veros != id(veros):
        solve.linear_solver = _get_solver(veros)
        solve.veros = id(veros)

    if veros.enable_cyclic_x:
        cyclic.setcyclic_x(sol)

    if boundary_val is None:
        boundary_val = sol

    z = np.prod(~veros.boundary_mask, axis=2).astype(np.bool)
    rhs[...] = np.where(z, rhs, boundary_val) # set right hand side on boundaries
    linear_solution = solve.linear_solver(rhs, sol)
    sol[...] = boundary_val
    sol[2:-2, 2:-2] = linear_solution.reshape(veros.nx + 4, veros.ny + 4)[2:-2, 2:-2]


solve.veros = None


@veros_method
def _get_solver(veros):
    matrix = _assemble_poisson_matrix(veros)
    if has_pyamg:
        preconditioner = _amg_preconditioner(veros, matrix)
    else:
        preconditioner = _jacobi_preconditioner(veros, matrix)
        # preconditioner.diagonal()[...] *= np.prod(~veros.boundary_mask, axis=2).astype(np.int).flatten()

    def scipy_solver(rhs, x0):
        solution, info = spalg.bicgstab(matrix, rhs.flatten(),
                                        x0=x0.flatten(), tol=veros.congr_epsilon * 1e-8,
                                        maxiter=veros.congr_max_iterations,
                                        M=preconditioner)
        if info > 0:
            warnings.warn("Streamfunction solver did not converge after {} iterations".format(info))
        return solution
    return scipy_solver


@veros_method
def _amg_preconditioner(veros, matrix):
    if veros.backend_name == "bohrium":
        near_null_space = np.ones((matrix.shape[0], 1), bohrium=False, dtype=veros.default_float_type)
    else:
        near_null_space = np.ones((matrix.shape[0], 1), dtype=veros.default_float_type)
    ml = pyamg.smoothed_aggregation_solver(matrix, near_null_space, np.copy(near_null_space), symmetry="nonsymmetric")
    return ml.aspreconditioner()


@veros_method
def _jacobi_preconditioner(veros, matrix):
    """
    Construct a simple Jacobi preconditioner
    """
    Z = np.ones((veros.nx + 4, veros.ny + 4), dtype=veros.default_float_type)
    Y = matrix.diagonal().copy().reshape(veros.nx + 4, veros.ny + 4)[2:-2, 2:-2]
    Z[2:-2, 2:-2] = np.where(Y != 0., 1. / Y, 1.)
    return scipy.sparse.dia_matrix((Z.flatten(), 0), shape=(Z.size, Z.size)).tocsr()


@veros_method
def _assemble_poisson_matrix(veros):
    """
    Construct a sparse matrix based on the stencil for the 2D Poisson equation.
    """
    # assemble diagonals
    main_diag = np.ones((veros.nx + 4, veros.ny + 4), dtype=veros.default_float_type)
{'fraction_non_alphanumeric': 0.1139481189732446, 'fraction_numerical': 0.012494906967268776, 'mean_word_length': 2.898358920063526, 'pattern_counts': {'":': 3, '<': 3, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6871344', 'n_tokens_mistral': 2598, 'n_tokens_neox': 2477, 'n_words': 566}