Test Problem Gallery

def poisson(grid, format='csr', dtype=float):
    """
    Generate discrete Poisson operator on regular grid.
    
    Creates finite difference discretization of -∇²u using
    standard 5-point (2D) or 7-point (3D) stencils.
    
    Parameters:
    - grid: tuple, grid dimensions:
        * (n,): 1D grid with n points  
        * (nx, ny): 2D grid with nx×ny points
        * (nx, ny, nz): 3D grid with nx×ny×nz points
    - format: str, sparse matrix format:
        * 'csr': compressed sparse row (default)
        * 'csc': compressed sparse column
        * 'coo': coordinate format
        * 'lil': list of lists format
    - dtype: data type, float or complex
    
    Returns:
    sparse matrix: Poisson operator matrix
    
    Raises:
    ValueError: if grid dimensions are invalid
    """

def gauge_laplacian(grid, beta=0.1, format='csr', dtype=float):
    """
    Generate gauge Laplacian operator.
    
    Creates Laplacian with gauge transformation, resulting in
    singular matrix with non-trivial nullspace.
    
    Parameters:
    - grid: tuple, grid dimensions
    - beta: float, gauge parameter
    - format: str, sparse matrix format  
    - dtype: data type
    
    Returns:
    sparse matrix: gauge Laplacian operator
    """

import pyamg
import numpy as np

# 1D Poisson problem
A1d = pyamg.gallery.poisson((100,))
print(f"1D Poisson: {A1d.shape}, {A1d.nnz} nonzeros")

# 2D Poisson problem  
A2d = pyamg.gallery.poisson((50, 50))
print(f"2D Poisson: {A2d.shape}, {A2d.nnz} nonzeros")

# 3D Poisson problem
A3d = pyamg.gallery.poisson((20, 20, 20))
print(f"3D Poisson: {A3d.shape}, {A3d.nnz} nonzeros")

# Different matrix format
A_coo = pyamg.gallery.poisson((40, 40), format='coo')

# Gauge Laplacian (singular)
A_gauge = pyamg.gallery.gauge_laplacian((30, 30))

def linear_elasticity(grid, format='csr', dtype=float):
    """
    Generate linear elasticity operator for displacement formulation.
    
    Creates finite element discretization of linear elasticity
    equations using Q1 elements. Results in block system with
    rigid body modes in nullspace.
    
    Parameters:
    - grid: tuple, grid dimensions (nx, ny) for 2D problems
    - format: str, sparse matrix format
    - dtype: data type
    
    Returns:
    sparse matrix: elasticity operator (2*nx*ny × 2*nx*ny for 2D)
    
    Notes:
    - Nullspace contains rigid body modes (translation + rotation)
    - Requires special treatment in AMG (near-nullspace)
    - Best solved with Smoothed Aggregation using rigid body modes
    """

def linear_elasticity_p1(vertices, elements, format='csr'):
    """
    Generate P1 linear elasticity operator on triangular mesh.
    
    Creates finite element discretization using P1 (linear)
    elements on unstructured triangular mesh.
    
    Parameters:
    - vertices: array, vertex coordinates (n_vertices × 2)
    - elements: array, element connectivity (n_elements × 3)  
    - format: str, sparse matrix format
    
    Returns:
    sparse matrix: P1 elasticity operator
    """

# 2D linear elasticity
A_elastic = pyamg.gallery.linear_elasticity((40, 40))
print(f"Elasticity: {A_elastic.shape}, {A_elastic.nnz} nonzeros")

# Solve with rigid body modes
B = np.ones((A_elastic.shape[0], 3))  # Translation + rotation modes
ml = pyamg.smoothed_aggregation_solver(A_elastic, B=B)

# P1 elasticity on triangle mesh
vertices = np.array([[0,0], [1,0], [1,1], [0,1]])
elements = np.array([[0,1,2], [0,2,3]])  
A_p1 = pyamg.gallery.linear_elasticity_p1(vertices, elements)

def stencil_grid(stencil, grid, format='csr', dtype=float):
    """
    Generate matrix from stencil on regular grid.
    
    Creates sparse matrix by applying given stencil to all
    interior points of regular grid with appropriate boundary conditions.
    
    Parameters:
    - stencil: array, stencil weights:
        * 1D: shape (2*r+1,) for radius r stencil
        * 2D: shape (2*r+1, 2*s+1) for radius (r,s) stencil  
        * 3D: shape (2*r+1, 2*s+1, 2*t+1) stencil
    - grid: tuple, grid dimensions
    - format: str, sparse matrix format
    - dtype: data type
    
    Returns:
    sparse matrix: stencil-generated operator
    
    Raises:
    ValueError: if stencil and grid dimensions are incompatible
    """

def diffusion_stencil_2d(epsilon=1.0, theta=0.0, type='FE'):
    """
    Generate 2D diffusion stencil with anisotropy.
    
    Creates stencil for anisotropic diffusion operator with
    specified anisotropy ratio and orientation.
    
    Parameters:
    - epsilon: float, anisotropy ratio (>1 for strong anisotropy)
    - theta: float, rotation angle in radians
    - type: str, discretization type:
        * 'FE': finite element  
        * 'FD': finite difference
    
    Returns:
    array: 3×3 stencil for 2D diffusion operator
    """

# Custom 1D stencil (second derivative)
stencil_1d = np.array([1, -2, 1])
A_custom = pyamg.gallery.stencil_grid(stencil_1d, (100,))

# 2D Laplacian stencil
stencil_2d = np.array([[0, 1, 0], 
                       [1, -4, 1], 
                       [0, 1, 0]])
A_laplace = pyamg.gallery.stencil_grid(stencil_2d, (50, 50))

# Anisotropic diffusion
stencil_aniso = pyamg.gallery.diffusion_stencil_2d(epsilon=100, theta=np.pi/4)
A_aniso = pyamg.gallery.stencil_grid(stencil_aniso, (40, 40))

def sprand(m, n, density, format='csr', dtype=float, 
           distribution='uniform', **kwargs):
    """
    Generate random sparse matrix.
    
    Creates sparse matrix with specified density and random
    value distribution, useful for testing and benchmarking.
    
    Parameters:
    - m: int, number of rows
    - n: int, number of columns  
    - density: float, fraction of nonzero entries (0 < density ≤ 1)
    - format: str, sparse matrix format
    - dtype: data type
    - distribution: str, random distribution:
        * 'uniform': uniform distribution [0,1]
        * 'normal': normal distribution N(0,1)
    - **kwargs: distribution-specific parameters
    
    Returns:
    sparse matrix: random sparse matrix
    """

def load_example(name):
    """
    Load example matrix from PyAMG collection.
    
    Loads predefined test matrices commonly used in
    multigrid research and benchmarking.
    
    Parameters:
    - name: str, example name:
        * 'airfoil': unstructured airfoil mesh
        * 'bar': structural mechanics bar
        * 'fem_2d': 2D finite element matrix
        * Additional examples available
    
    Returns:
    sparse matrix: loaded example matrix
    
    Raises:
    ValueError: if example name is not recognized
    """

# Random sparse matrix
A_rand = pyamg.gallery.sprand(1000, 1000, density=0.01)

# Random with normal distribution  
A_normal = pyamg.gallery.sprand(500, 500, density=0.02, 
                               distribution='normal')

# Load example problem
try:
    A_example = pyamg.gallery.load_example('airfoil')
    print(f"Airfoil: {A_example.shape}, {A_example.nnz} nonzeros")
except ValueError:
    print("Example not available")

class Mesh:
    """
    Mesh representation for finite element problems.
    
    Manages vertices, elements, and connectivity for
    unstructured mesh-based discretizations.
    
    Attributes:
    - vertices: array of vertex coordinates
    - elements: array of element connectivity
    - boundary: boundary node information
    """
    
    def __init__(self, vertices, elements):
        """
        Initialize mesh from vertices and elements.
        
        Parameters:
        - vertices: array, vertex coordinates (n_vertices × dim)
        - elements: array, element connectivity (n_elements × nodes_per_element)
        """

def regular_triangle_mesh(nx, ny):
    """
    Generate regular triangular mesh.
    
    Creates structured triangular mesh on unit square
    with specified resolution.
    
    Parameters:
    - nx: int, number of divisions in x direction
    - ny: int, number of divisions in y direction
    
    Returns:
    tuple: (vertices, elements) arrays for triangle mesh
    """

def stokes(mesh, format='csr'):
    """
    Generate Stokes problem operator.
    
    Creates mixed finite element discretization of
    Stokes equations (incompressible flow).
    
    Parameters:
    - mesh: Mesh object or tuple (vertices, elements)
    - format: str, sparse matrix format
    
    Returns:
    sparse matrix: Stokes operator (saddle point system)
    """

# Create triangular mesh
vertices, elements = pyamg.gallery.regular_triangle_mesh(20, 20)
mesh = pyamg.gallery.Mesh(vertices, elements)

# Generate Stokes problem
A_stokes = pyamg.gallery.stokes(mesh)
print(f"Stokes: {A_stokes.shape}, {A_stokes.nnz} nonzeros")

def demo(**kwargs):
    """
    Run PyAMG demonstration showing solver capabilities.
    
    Interactive demo that creates test problems, constructs
    AMG solvers, and displays convergence behavior and
    hierarchy information.
    
    Parameters:
    - **kwargs: demo configuration options
    
    Returns:
    None (prints demo output)
    """

# Run basic demo
pyamg.gallery.demo()

# Demo with custom parameters  
pyamg.gallery.demo(grid_size=(100, 100), solver='smoothed_aggregation')