SciCode++
Switch benchSciCode++
Request dataset samplesby

Sketched Block Krylov QR for a Stiff Chebyshev Collocation SystemMathematics

13 verifiers

TaxonomyMathematicsNumerical Linear Algebra

Task prompt

  • Gemini 3.1 Pro Preview75%of 8 runs
  • GPT-5.6 (sol)37.5%of 8 runs

Prompt

Calculate the sketched block-GMRES approximation of a stiff Chebyshev-Lobatto collocation linear system while certifying the coupled QR, Krylov, and residual invariants of the randomized block basis. Let n≥4n \ge 4 and let the Chebyshev-Lobatto nodes be xj=cos⁡(πj/(n−1))x_j=\cos(\pi j/(n-1)) for j=0,…,n−1j=0,\ldots,n-1, ordered from 11 to −1-1. Let DD denote the unique nodal differentiation matrix that maps the nodal values of any polynomial of degree at most n−1n-1 to the nodal values of its derivative on these nodes. Define D(2)=DDD^{(2)}=D D. For parameters ν>0\nu>0, a∈Ra\in\mathbb{R}, and ρ≥0\rho\ge 0, define the collocation matrix A∈Rn×nA\in\mathbb{R}^{n\times n} by the interior differential operator −νD(2)+aD+ρI-\nu D^{(2)}+aD+\rho I on rows 1,…,n−21,\ldots,n-2, with Dirichlet boundary rows A0,:=e0⊤A_{0,:}=e_0^\top and An−1,:=en−1⊤A_{n-1,:}=e_{n-1}^\top. The inputs B,X0∈Rn×sB,X_0\in\mathbb{R}^{n\times s} contain s≥1s\ge 1 right-hand sides and initial guesses. Define the initial residual block R0=B−AX0R_0=B-AX_0. For an integer p≥1p\ge 1, define the block Krylov matrix W=[R0,AR0,A2R0,…,Ap−1R0]∈Rn×m,m=ps. W=[R_0,AR_0,A^2R_0,\ldots,A^{p-1}R_0]\in\mathbb{R}^{n\times m},\qquad m=ps. The admissible inputs are those for which WW has full column rank, Θ\Theta satisfies k≥mk\geq m, the required sketched block Gram matrices are positive definite, and σmin⁡(ΘAQ)>εmachσmax⁡(ΘAQ)\sigma_{\min}(\Theta A Q)>\varepsilon_{\mathrm{mach}}\sigma_{\max}(\Theta A Q). The returned Q,RQ,R are defined by the deterministic sketched block Gram-Schmidt recurrence over the pp consecutive blocks of width ss: each block residual is orthogonalized against all previous returned blocks through the sketched Euclidean inner product induced by Θ\Theta, each diagonal block is normalized by the upper Cholesky factor of its sketched Gram matrix with positive diagonal, and a final full-space CholeskyQR cleanup is applied so that Q⊤QQ^\top Q is orthonormal to fine precision while preserving W=QRW=QR. The final RR must be upper triangular with strictly positive diagonal. The returned approximation XX is the unique element of X0+range⁡(Q)X_0+\operatorname{range}(Q) whose coefficient matrix Y∈Rm×sY\in\mathbb{R}^{m\times s} is the minimum-Frobenius-norm minimizer of min⁡Y∈Rm×s∥Θ(A(X0+QY)−B)∥F2. \min_{Y\in\mathbb{R}^{m\times s}}\|\Theta(A(X_0+QY)-B)\|_F^2. Let Z=ΘAQZ=\Theta A Q and rΘ=Θ(AX−B)r_\Theta=\Theta(AX-B). Define \texttt{normal_equation_defect} = \frac{\|Z^\top r_\Theta\|_F} {\max(\|Z\|_F\|r_\Theta\|_F,\varepsilon_{\mathrm{mach}})}. For p>1p>1, let Q−=Q(:,1:m−s)Q_-=Q_{(:,1:m-s)}, H=Q⊤AQ−H=Q^\top A Q_-, and let HfarH_{\mathrm{far}} contain the blocks of HH below the first block subdiagonal and zeros elsewhere. Define \texttt{arnoldi_defect} = \frac{\|AQ_- - QH\|_F} {\max(\|AQ_-\|_F,\varepsilon_{\mathrm{mach}})}, \qquad \texttt{hessenberg_defect} = \frac{\|H_{\mathrm{far}}\|_F} {\max(\|H\|_F,\varepsilon_{\mathrm{mach}})}. For p=1p=1, set the last two defects to zero. Correctness also requires W≈QRW\approx QR, Q⊤Q≈ImQ^\top Q\approx I_m, and an upper-triangular RR with positive diagonal. If any admissibility condition is violated, raise ValueError. Return a flat dictionary with the keys and meanings stated in the function declaration. Norm diagnostics are dimensionless relative quantities except for residual norms, which have the same units as the entries of BB.

def solve_coupled_chebyshev_sketch_gmres(n, B, X0, theta, nu, adv, rho, krylov_blocks):
    '''
    Inputs:
        n : int, number of Chebyshev-Lobatto nodes, n >= 4
        B : np.ndarray, shape (n, s), right-hand side block
        X0 : np.ndarray, shape (n, s), initial guess block
        theta : np.ndarray, shape (k, n), real sketching matrix with k >= krylov_blocks*s
        nu : float, positive diffusion coefficient
        adv : float, advection coefficient
        rho : float, nonnegative reaction coefficient
        krylov_blocks : int, number of block Krylov powers, at least 1
    Outputs:
        result : dict
            'X' : np.ndarray, shape (n, s), sketched minimum-residual approximation
            'Q' : np.ndarray, shape (n, krylov_blocks*s), cleaned block Krylov basis
            'R' : np.ndarray, shape (krylov_blocks*s, krylov_blocks*s), upper triangular factor
            'Y' : np.ndarray, shape (krylov_blocks*s, s), coefficient matrix in the basis Q
            'krylov_matrix' : np.ndarray, shape (n, krylov_blocks*s), block Krylov matrix W
            'qr_relative_residual' : float, relative Frobenius defect of W = Q R
            'orthogonality_defect' : float, spectral norm of Q.T Q - I
            'triangular_defect' : float, relative Frobenius norm of the strict lower triangle of R
            'positive_diagonal' : bool, whether all diagonal entries of R are positive
            'sketched_residual_norm' : float, Frobenius norm of theta @ (A @ X - B)
            'true_residual_norm' : float, Frobenius norm of A @ X - B
            'normal_equation_defect' : float, relative sketched Petrov-Galerkin defect
            'arnoldi_defect' : float, relative defect of A Q_previous lying in range(Q)
            'hessenberg_defect' : float, relative violation of block upper-Hessenberg structure
    '''
    return result

Background

This problem couples three numerical-linear-algebra structures. The first is a pseudospectral discretization on the nonuniform Chebyshev-Lobatto grid, where clustered endpoint nodes make differentiation matrices highly nonnormal and increasingly ill-conditioned as nn grows. The second is a randomized block QR factorization of the Krylov matrix W=[R0,AR0,…,Ap−1R0]W=[R_0,AR_0,\ldots,A^{p-1}R_0], where successive block powers can become nearly linearly dependent when the spectrum of AA is clustered or when diffusion and reaction create separated scales. The third is a sketched residual minimization problem in the computed Krylov subspace. The central invariants are algebraic rather than statistical: the factorization invariant is W=QRW=QR with RR upper triangular and rjj>0r_{jj}>0; the basis invariant is Q⊤Q=IQ^\top Q=I up to numerical tolerance; and the residual invariant is the sketched Petrov-Galerkin condition (ΘAQ)⊤Θ(AX−B)=0(\Theta A Q)^\top\Theta(A X-B)=0 for the returned minimizer. The Chebyshev operator introduces stiffness through the second derivative matrix, while the randomized sketch couples the high-dimensional basis construction to a lower-dimensional inner product. A solution that satisfies only one of these structures will generally violate the others: for example, a basis with poor orthogonality corrupts the residual minimization, and a residual minimizer outside the Krylov range violates the range-preservation constraint.

Parameters and returns

Inputs: n : int, number of Chebyshev-Lobatto nodes, n >= 4 B : np.ndarray, shape (n, s), right-hand side block X0 : np.ndarray, shape (n, s), initial guess block theta : np.ndarray, shape (k, n), real sketching matrix with k >= krylov_blockss nu : float, positive diffusion coefficient adv : float, advection coefficient rho : float, nonnegative reaction coefficient krylov_blocks : int, number of block Krylov powers, at least 1 Outputs: result : dict 'X' : np.ndarray, shape (n, s), sketched minimum-residual approximation 'Q' : np.ndarray, shape (n, krylov_blockss), cleaned block Krylov basis 'R' : np.ndarray, shape (krylov_blockss, krylov_blockss), upper triangular factor 'Y' : np.ndarray, shape (krylov_blockss, s), coefficient matrix in the basis Q 'krylov_matrix' : np.ndarray, shape (n, krylov_blockss), block Krylov matrix W 'qr_relative_residual' : float, relative Frobenius defect of W = Q R 'orthogonality_defect' : float, spectral norm of Q.T Q - I 'triangular_defect' : float, relative Frobenius norm of the strict lower triangle of R 'positive_diagonal' : bool, whether all diagonal entries of R are positive 'sketched_residual_norm' : float, Frobenius norm of theta @ (A @ X - B) 'true_residual_norm' : float, Frobenius norm of A @ X - B 'normal_equation_defect' : float, relative sketched Petrov-Galerkin defect 'arnoldi_defect' : float, relative defect of A Q_previous lying in range(Q) 'hessenberg_defect' : float, relative violation of block upper-Hessenberg structure

Required dependencies

import numpy as np  # numpy==2.4.4
import scipy.linalg as la  # scipy==1.17.1

Published by the benchmark's authors. Opening these shows how the problem is solved and what the graders assert.

Test case 1 of 5
def test_case_1():
    import numpy as np
    import scipy.linalg as la
    # Direct active-regime oracle for coupled Chebyshev Krylov QR and sketched GMRES; catches arith_swap in A, W, QR, or residual formation.

    def Aop(n, nu, adv, rho):
        j = np.arange(n, dtype=float); x = np.cos(np.pi * j / float(n - 1)); c = np.ones(n); c[0] = 2.0; c[-1] = 2.0; c = c * ((-1.0) ** np.arange(n)); dX = x[:, None] - x[None, :]; D = np.zeros((n, n)); mask = ~np.eye(n, dtype=bool); C = c[:, None] / c[None, :]; D[mask] = C[mask] / dX[mask]; D[np.diag_indices(n)] = -np.sum(D, axis=1); A = -nu * (D @ D) + adv * D + rho * np.eye(n); A[0, :] = 0.0; A[0, 0] = 1.0; A[-1, :] = 0.0; A[-1, -1] = 1.0; return A
    n = 8; nu = 0.06; adv = 1.9; rho = 0.7; p = 3
    B = np.array([[1.2, -0.7], [2.5, 0.3], [-1.1, 1.8], [0.6, -2.2], [3.0, 1.1], [-0.9, 2.4], [1.7, -1.5], [-2.0, 0.8]])
    X0 = np.array([[0.2, 0.5], [-0.4, 0.1], [0.7, -0.3], [-0.6, 0.9], [0.3, -0.8], [1.0, 0.4], [-0.2, -0.5], [0.8, 0.6]])
    theta = np.eye(n); theta[1, 0] = 0.22; theta[4, 2] = -0.35; theta[7, 5] = 0.18
    result = solve_coupled_chebyshev_sketch_gmres(n, B, X0, theta, nu, adv, rho, p)
    A = Aop(n, nu, adv, rho); R0 = B - A @ X0; blocks = []; V = R0.copy()
    for _ in range(p):
        blocks.append(V.copy()); V = A @ V
    W = np.hstack(blocks); q, r = la.qr(W, mode='economic'); sgn = np.where(np.diag(r) < 0.0, -1.0, 1.0); q = q * sgn; r = sgn[:, None] * r
    Z = theta @ (A @ q); S = theta @ (B - A @ X0); Yexp, _, _, _ = la.lstsq(Z, S, lapack_driver='gelsd'); Xexp = X0 + q @ Yexp
    assert np.allclose(result['X'], Xexp, rtol=1e-8, atol=1e-8)
    assert np.allclose(result['R'], r, rtol=1e-8, atol=1e-8)
    assert abs(result['sketched_residual_norm'] - la.norm(theta @ (A @ Xexp - B), ord='fro')) < 1e-7
    expected_tri = la.norm(np.tril(result['R'], -1), ord='fro') / max(
        la.norm(result['R'], ord='fro'), np.finfo(float).eps
    )
    assert np.isclose(
        result['triangular_defect'], expected_tri, rtol=1e-12, atol=1e-12
    )
    n_bad = 4
    B_bad = np.zeros((n_bad, 1))
    B_bad[0, 0] = 1.0
    X0_bad = np.zeros_like(B_bad)
    A_bad = Aop(n_bad, 0.1, 0.2, 0.3)
    theta_bad = np.array([[A_bad[1, 0], -1.0, 0.0, 0.0]])

    try:
        solve_coupled_chebyshev_sketch_gmres(
            n_bad, B_bad, X0_bad, theta_bad, 0.1, 0.2, 0.3, 1
        )
    except ValueError:
        pass
    else:
        raise AssertionError('rank-deficient theta @ A @ Q must raise ValueError')
Test case 2 of 5
def test_case_2():
    import numpy as np
    import scipy.linalg as la
    # Active sign/factor sensitivity for -nu D2 + adv D + rho I with all terms nonzero; catches +/-, */ mutations and dropping cos in Chebyshev nodes.

    def Aop(n, nu, adv, rho):
        j = np.arange(n, dtype=float); x = np.cos(np.pi * j / float(n - 1)); c = np.ones(n); c[0] = 2.0; c[-1] = 2.0; c = c * ((-1.0) ** np.arange(n)); dX = x[:, None] - x[None, :]; D = np.zeros((n, n)); mask = ~np.eye(n, dtype=bool); C = c[:, None] / c[None, :]; D[mask] = C[mask] / dX[mask]; D[np.diag_indices(n)] = -np.sum(D, axis=1); A = -nu * (D @ D) + adv * D + rho * np.eye(n); A[0, :] = 0.0; A[0, 0] = 1.0; A[-1, :] = 0.0; A[-1, -1] = 1.0; return A
    n = 9; nu = 0.075; adv = -2.75; rho = 1.3; p = 2
    B = np.reshape(np.linspace(-2.3, 3.1, 18), (9, 2))
    X0 = np.reshape(np.linspace(1.4, -1.2, 18), (9, 2))
    theta = np.vstack([np.eye(n)[:8], np.array([[0.3, -0.2, 0.5, 0.1, -0.4, 0.2, 0.0, 0.6, -0.1]])])
    result = solve_coupled_chebyshev_sketch_gmres(n, B, X0, theta, nu, adv, rho, p)
    A = Aop(n, nu, adv, rho); R0 = B - A @ X0; W = np.hstack([R0, A @ R0]); q, r = la.qr(W, mode='economic'); sgn = np.where(np.diag(r) < 0.0, -1.0, 1.0); q = q * sgn
    Yexp, _, _, _ = la.lstsq(theta @ (A @ q), theta @ (B - A @ X0), lapack_driver='gelsd'); Xexp = X0 + q @ Yexp
    expected_true = la.norm(A @ Xexp - B, ord='fro')
    assert np.allclose(result['krylov_matrix'], W, rtol=1e-9, atol=1e-7)
    assert abs(result['true_residual_norm'] - expected_true) < 1e-7
    assert result['true_residual_norm'] > 1e-4
    assert result['normal_equation_defect'] < 1e-10
Test case 3 of 5
def test_case_3():
    import numpy as np
    # Boundary/threshold discrimination: theta has exactly k=m rows and rho is positive but near the rho>=0 boundary; catches < versus <= and off-by-one k checks.

    def Aop(n, nu, adv, rho):
        j = np.arange(n, dtype=float); x = np.cos(np.pi * j / float(n - 1)); c = np.ones(n); c[0] = 2.0; c[-1] = 2.0; c = c * ((-1.0) ** np.arange(n)); dX = x[:, None] - x[None, :]; D = np.zeros((n, n)); mask = ~np.eye(n, dtype=bool); C = c[:, None] / c[None, :]; D[mask] = C[mask] / dX[mask]; D[np.diag_indices(n)] = -np.sum(D, axis=1); A = -nu * (D @ D) + adv * D + rho * np.eye(n); A[0, :] = 0.0; A[0, 0] = 1.0; A[-1, :] = 0.0; A[-1, -1] = 1.0; return A
    n = 7; nu = 0.09; adv = 1.4; rho = 1e-9; p = 2
    B = np.array([[1.0, 0.4], [-0.8, 1.7], [2.2, -1.1], [0.5, 2.5], [-1.9, 0.8], [1.3, -0.6], [0.7, 1.1]])
    X0 = np.array([[0.3, -0.2], [0.6, 0.1], [-0.5, 0.9], [1.0, -0.7], [0.2, 0.4], [-0.8, 0.5], [0.9, -0.3]])
    theta = np.array([[1.0, 0.2, -0.1, 0.3, 0.0, -0.2, 0.4], [0.1, 1.2, 0.3, -0.2, 0.5, 0.0, -0.1], [-0.3, 0.4, 1.1, 0.2, -0.1, 0.6, 0.0], [0.2, -0.1, 0.5, 1.3, 0.4, -0.3, 0.2]])
    result = solve_coupled_chebyshev_sketch_gmres(n, B, X0, theta, nu, adv, rho, p)
    A = Aop(n, nu, adv, rho); R0 = B - A @ X0; W = np.hstack([R0, A @ R0])
    assert np.allclose(result['krylov_matrix'], W, rtol=1e-8, atol=1e-8)
    assert np.all(np.diag(result['R']) > 0.0)
    assert abs(result['qr_relative_residual']) < 1e-10
Test case 4 of 5
def test_case_4():
    import numpy as np
    import scipy.linalg as la
    # Shape/orientation sensitivity with B 9x3, theta 8x9, and block width 3; catches removing .T in Gram products or swapping axes.

    def Aop(n, nu, adv, rho):
        j = np.arange(n, dtype=float); x = np.cos(np.pi * j / float(n - 1)); c = np.ones(n); c[0] = 2.0; c[-1] = 2.0; c = c * ((-1.0) ** np.arange(n)); dX = x[:, None] - x[None, :]; D = np.zeros((n, n)); mask = ~np.eye(n, dtype=bool); C = c[:, None] / c[None, :]; D[mask] = C[mask] / dX[mask]; D[np.diag_indices(n)] = -np.sum(D, axis=1); A = -nu * (D @ D) + adv * D + rho * np.eye(n); A[0, :] = 0.0; A[0, 0] = 1.0; A[-1, :] = 0.0; A[-1, -1] = 1.0; return A
    n = 9; nu = 0.055; adv = 2.1; rho = 0.8; p = 2
    B = np.reshape(np.array([1.0, -0.5, 2.0, 0.7, 1.3, -1.1, -0.8, 2.4, 0.6, 1.5, -2.0, 0.9, 2.2, 0.4, -1.7, -1.3, 1.8, 0.2, 0.5, -0.9, 2.7, 1.1, 0.3, -2.5, -0.4, 2.1, 1.6]), (9, 3))
    X0 = np.reshape(np.linspace(-0.9, 1.1, 27), (9, 3))
    theta = np.array([[1.0, 0.2, 0.0, -0.1, 0.3, 0.0, 0.1, -0.2, 0.4], [0.0, 1.3, -0.2, 0.2, 0.0, 0.4, -0.1, 0.3, -0.3], [0.2, -0.1, 1.1, 0.0, -0.3, 0.2, 0.0, 0.1, 0.5], [-0.1, 0.3, 0.2, 1.4, 0.1, -0.2, 0.4, 0.0, -0.2], [0.3, 0.0, -0.4, 0.1, 1.2, 0.2, -0.2, 0.5, 0.1], [0.0, 0.2, 0.1, -0.3, 0.4, 1.5, 0.3, -0.1, 0.2], [0.4, -0.2, 0.3, 0.0, -0.1, 0.2, 1.1, 0.3, -0.4], [-0.2, 0.1, 0.0, 0.5, 0.2, -0.3, 0.4, 1.2, 0.6]])
    result = solve_coupled_chebyshev_sketch_gmres(n, B, X0, theta, nu, adv, rho, p)
    A = Aop(n, nu, adv, rho); R0 = B - A @ X0; W = np.hstack([R0, A @ R0]); q, r = la.qr(W, mode='economic'); sgn = np.where(np.diag(r) < 0.0, -1.0, 1.0); q = q * sgn
    Yexp, _, _, _ = la.lstsq(theta @ (A @ q), theta @ (B - A @ X0), lapack_driver='gelsd')
    assert np.allclose(result['Y'], Yexp, rtol=1e-8, atol=1e-8)
    assert abs(result['orthogonality_defect']) < 1e-10
    assert result['hessenberg_defect'] < 1e-10
    assert result['arnoldi_defect'] < 1e-10
Test case 5 of 5
def test_case_5():
    import numpy as np
    import scipy.linalg as la
    # Iteration/index precision for four nontrivial Krylov powers [R0,AR0,A^2R0,A^3R0]; catches off-by-one loop bounds in Krylov construction and diagnostics.

    def Aop(n, nu, adv, rho):
        j = np.arange(n, dtype=float); x = np.cos(np.pi * j / float(n - 1)); c = np.ones(n); c[0] = 2.0; c[-1] = 2.0; c = c * ((-1.0) ** np.arange(n)); dX = x[:, None] - x[None, :]; D = np.zeros((n, n)); mask = ~np.eye(n, dtype=bool); C = c[:, None] / c[None, :]; D[mask] = C[mask] / dX[mask]; D[np.diag_indices(n)] = -np.sum(D, axis=1); A = -nu * (D @ D) + adv * D + rho * np.eye(n); A[0, :] = 0.0; A[0, 0] = 1.0; A[-1, :] = 0.0; A[-1, -1] = 1.0; return A
    n = 8; nu = 0.045; adv = 1.6; rho = 0.95; p = 4
    B = np.array([[1.0], [-1.4], [2.2], [0.6], [-0.9], [1.8], [-2.1], [0.7]])
    X0 = np.array([[0.5], [-0.2], [0.9], [-0.7], [0.3], [1.1], [-0.4], [0.8]])
    theta = np.eye(n); theta[0, 2] = 0.28; theta[3, 6] = -0.33; theta[5, 1] = 0.19
    result = solve_coupled_chebyshev_sketch_gmres(n, B, X0, theta, nu, adv, rho, p)
    A = Aop(n, nu, adv, rho); R0 = B - A @ X0; blocks = []; V = R0.copy()
    for _ in range(p):
        blocks.append(V.copy()); V = A @ V
    W = np.hstack(blocks)
    q, r = la.qr(W, mode='economic'); sgn = np.where(np.diag(r) < 0.0, -1.0, 1.0); r = sgn[:, None] * r
    assert np.allclose(result['krylov_matrix'], W, rtol=1e-9, atol=1e-6)
    assert np.allclose(result['R'], r, rtol=1e-8, atol=1e-6)
    assert la.norm(result['krylov_matrix'][:, -1], ord=2) > 1e3
    assert result['arnoldi_defect'] < 1e-10
Solution
import numpy as np
import scipy.linalg as la


def _norm_fro(X):
    X = np.asarray(X)
    if X.size == 0:
        return 0.0
    return float(la.norm(X, ord='fro'))


def _norm_2(X):
    X = np.asarray(X)
    if X.size == 0:
        return 0.0
    return float(la.norm(X, ord=2))


def _chol_upper_spd(G, name):
    G = np.asarray(G, dtype=float)
    if G.ndim != 2 or G.shape[0] != G.shape[1]:
        raise ValueError(name + ' must be square')
    G = 0.5 * (G + G.T)
    if not np.all(np.isfinite(G)):
        raise ValueError(name + ' contains nonfinite entries')
    try:
        return la.cholesky(G, lower=False, check_finite=True)
    except la.LinAlgError as exc:
        raise ValueError(name + ' is not positive definite') from exc


def _right_solve_upper(V, R):
    if not np.all(np.isfinite(V)) or not np.all(np.isfinite(R)):
        raise ValueError('nonfinite triangular solve input')
    diagonal = np.abs(np.diag(R))
    scale = max(float(la.norm(R, ord=np.inf)), 1.0)
    tolerance = np.finfo(float).eps * scale

    if np.any(diagonal <= tolerance):
        raise ValueError('singular triangular factor')
    return la.solve_triangular(R, V.T, trans='T', lower=False, check_finite=True).T


def rbgs_block_cholesky_qr(W, block_width, theta):
    W = np.asarray(W, dtype=float)
    theta = np.asarray(theta, dtype=float)
    eps = np.finfo(float).eps
    if W.ndim != 2 or theta.ndim != 2:
        raise ValueError('W and theta must be two-dimensional')
    n, m = W.shape
    k, nt = theta.shape
    if nt != n:
        raise ValueError('theta has incompatible column dimension')
    if not isinstance(block_width, (int, np.integer)) or block_width <= 0:
        raise ValueError('block_width must be a positive integer')
    b = int(block_width)
    if m == 0 or m % b != 0:
        raise ValueError('number of columns of W must be positive and divisible by block_width')
    if k < m:
        raise ValueError('theta must have at least as many rows as W has columns')
    if not np.all(np.isfinite(W)) or not np.all(np.isfinite(theta)):
        raise ValueError('inputs contain nonfinite entries')
    p = m // b
    Q = np.zeros((n, m), dtype=float)
    R = np.zeros((m, m), dtype=float)
    thetaQ = np.zeros((k, m), dtype=float)
    for i in range(p):
        cs = slice(i * b, (i + 1) * b)
        V = W[:, cs].copy()
        TV = theta @ V
        for j in range(i):
            rs = slice(j * b, (j + 1) * b)
            C = thetaQ[:, rs].T @ TV
            R[rs, cs] = C
            V = V - Q[:, rs] @ C
            TV = TV - thetaQ[:, rs] @ C
        G = TV.T @ TV
        Rii = _chol_upper_spd(G, 'sketched block Gram matrix')
        R[cs, cs] = Rii
        Q[:, cs] = _right_solve_upper(V, Rii)
        thetaQ[:, cs] = _right_solve_upper(TV, Rii)
    Gfull = Q.T @ Q
    Cfull = _chol_upper_spd(Gfull, 'full cleanup Gram matrix')
    Q = _right_solve_upper(Q, Cfull)
    R = Cfull @ R
    thetaQ = theta @ Q
    rel = _norm_fro(W - Q @ R) / max(_norm_fro(W), eps)
    orth = _norm_2(Q.T @ Q - np.eye(m))
    sketch_orth = _norm_2(thetaQ.T @ thetaQ - np.eye(m))
    tri = _norm_fro(np.tril(R, -1)) / max(_norm_fro(R), eps)
    mind = float(np.min(np.diag(R))) if m else 0.0
    result = {
        'Q': Q,
        'R': R,
        'thetaQ': thetaQ,
        'relative_residual': float(rel),
        'orthogonality_defect': float(orth),
        'sketch_orthogonality_defect': float(sketch_orth),
        'triangular_defect': float(tri),
        'min_diagonal_R': mind,
        'rank_ok': bool(mind > 0.0)
    }
    return result


def _chebyshev_lobatto_operator(n, nu, adv, rho):
    if not isinstance(n, (int, np.integer)) or int(n) < 4:
        raise ValueError('n must be an integer at least 4')
    n = int(n)
    nu = float(nu)
    adv = float(adv)
    rho = float(rho)
    if not np.isfinite(nu) or not np.isfinite(adv) or not np.isfinite(rho):
        raise ValueError('operator parameters must be finite')
    if nu <= 0.0 or rho < 0.0:
        raise ValueError('nu must be positive and rho must be nonnegative')
    j = np.arange(n, dtype=float)
    x = np.cos(np.pi * j / float(n - 1))
    c = np.ones(n, dtype=float)
    c[0] = 2.0
    c[-1] = 2.0
    c = c * ((-1.0) ** np.arange(n))
    X = x[:, None]
    dX = X - X.T
    D = np.zeros((n, n), dtype=float)
    mask = ~np.eye(n, dtype=bool)
    C = c[:, None] / c[None, :]
    D[mask] = C[mask] / dX[mask]
    D[np.diag_indices(n)] = -np.sum(D, axis=1)
    D2 = D @ D
    A = -nu * D2 + adv * D + rho * np.eye(n)
    A[0, :] = 0.0
    A[0, 0] = 1.0
    A[-1, :] = 0.0
    A[-1, -1] = 1.0
    if not np.all(np.isfinite(A)):
        raise ValueError('constructed operator is nonfinite')
    return A


def chebyshev_sketched_trial_gmres(n, B, X0, Q, theta, nu, adv, rho):
    B = np.asarray(B, dtype=float)
    X0 = np.asarray(X0, dtype=float)
    Q = np.asarray(Q, dtype=float)
    theta = np.asarray(theta, dtype=float)
    eps = np.finfo(float).eps
    if B.ndim != 2 or X0.ndim != 2 or Q.ndim != 2 or theta.ndim != 2:
        raise ValueError('B, X0, Q, and theta must be two-dimensional')
    if B.shape != X0.shape:
        raise ValueError('B and X0 must have the same shape')
    if B.shape[0] != int(n) or Q.shape[0] != int(n) or theta.shape[1] != int(n):
        raise ValueError('incompatible n dimension')
    if not np.all(np.isfinite(B)) or not np.all(np.isfinite(X0)) or not np.all(np.isfinite(Q)) or not np.all(np.isfinite(theta)):
        raise ValueError('inputs contain nonfinite entries')
    A = _chebyshev_lobatto_operator(int(n), nu, adv, rho)
    Z = theta @ (A @ Q)
    S = theta @ (B - A @ X0)
    if not np.all(np.isfinite(Z)) or not np.all(np.isfinite(S)):
        raise ValueError('least-squares matrix or right-hand side is nonfinite')
    Y, _, _, _ = la.lstsq(
        Z, S, cond=eps, lapack_driver='gelsd', check_finite=True
    )
    X = X0 + Q @ Y
    residual = A @ X - B
    sk_res = theta @ residual
    normal = Z.T @ sk_res
    normal_denom = max(_norm_fro(Z) * _norm_fro(sk_res), eps)
    svals = la.svdvals(A)
    if svals.size == 0 or not np.all(np.isfinite(svals)):
        raise ValueError('operator singular values are invalid')
    cond = float(np.max(svals) / max(np.min(svals), eps))
    result = {
        'X': X,
        'Y': Y,
        'sketched_residual_norm': float(_norm_fro(sk_res)),
        'true_residual_norm': float(_norm_fro(residual)),
        'normal_equation_defect': float(_norm_fro(normal) / normal_denom),
        'operator_condition': cond
    }
    return result


def _block_krylov_matrix(A, R0, p):
    if not isinstance(p, (int, np.integer)) or int(p) < 1:
        raise ValueError('krylov_blocks must be a positive integer')
    p = int(p)
    blocks = []
    V = np.asarray(R0, dtype=float).copy()
    for _ in range(p):
        blocks.append(V.copy())
        V = A @ V
    W = np.hstack(blocks)
    if not np.all(np.isfinite(W)):
        raise ValueError('Krylov matrix contains nonfinite entries')
    return W


def _arnoldi_hessenberg_diagnostics(A, Q, block_width, p):
    eps = np.finfo(float).eps
    m = Q.shape[1]
    b = int(block_width)
    if p <= 1 or m <= b:
        return 0.0, 0.0
    Qprev = Q[:, :m - b]
    AQ = A @ Qprev
    H = Q.T @ AQ
    arn = _norm_fro(AQ - Q @ H) / max(_norm_fro(AQ), eps)
    Hbad = H.copy()
    for col_block in range(p - 1):
        cutoff = min((col_block + 2) * b, m)
        c0 = col_block * b
        c1 = (col_block + 1) * b
        if cutoff < m:
            Hbad[cutoff:, c0:c1] = H[cutoff:, c0:c1]
        Hbad[:cutoff, c0:c1] = 0.0
    hes = _norm_fro(Hbad) / max(_norm_fro(H), eps)
    return float(arn), float(hes)


def solve_coupled_chebyshev_sketch_gmres(n, B, X0, theta, nu, adv, rho, krylov_blocks):
    B = np.asarray(B, dtype=float)
    X0 = np.asarray(X0, dtype=float)
    theta = np.asarray(theta, dtype=float)
    eps = np.finfo(float).eps
    if B.ndim != 2 or X0.ndim != 2 or theta.ndim != 2:
        raise ValueError('B, X0, and theta must be two-dimensional')
    if B.shape != X0.shape:
        raise ValueError('B and X0 must have identical shape')
    if B.shape[0] != int(n) or theta.shape[1] != int(n):
        raise ValueError('incompatible n dimension')
    if not np.all(np.isfinite(B)) or not np.all(np.isfinite(X0)) or not np.all(np.isfinite(theta)):
        raise ValueError('inputs contain nonfinite entries')
    if not isinstance(krylov_blocks, (int, np.integer)) or int(krylov_blocks) < 1:
        raise ValueError('krylov_blocks must be a positive integer')
    p = int(krylov_blocks)
    s = B.shape[1]
    if s < 1:
        raise ValueError('there must be at least one right-hand side')
    m = p * s
    if theta.shape[0] < m:
        raise ValueError('theta has too few rows for the Krylov dimension')
    A = _chebyshev_lobatto_operator(int(n), nu, adv, rho)
    R0 = B - A @ X0
    W = _block_krylov_matrix(A, R0, p)
    qr = rbgs_block_cholesky_qr(W, s, theta)
    Q = qr['Q']
    R = qr['R']
    Z_main = theta @ (A @ Q)
    svals_main = la.svdvals(Z_main)
    if svals_main.size < m or svals_main[-1] <= eps * svals_main[0]:
        raise ValueError('theta @ A @ Q must have full numerical column rank')
    gm = chebyshev_sketched_trial_gmres(
        int(n), B, X0, Q, theta, nu, adv, rho
    )
    X = gm['X']
    Y = gm['Y']
    arnoldi_defect, hessenberg_defect = _arnoldi_hessenberg_diagnostics(A, Q, s, p)
    result = {
        'X': X,
        'Q': Q,
        'R': R,
        'Y': Y,
        'krylov_matrix': W,
        'qr_relative_residual': float(qr['relative_residual']),
        'orthogonality_defect': float(qr['orthogonality_defect']),
        'triangular_defect': float(qr['triangular_defect']),
        'positive_diagonal': bool(np.all(np.diag(R) > 0.0)),
        'sketched_residual_norm': float(gm['sketched_residual_norm']),
        'true_residual_norm': float(gm['true_residual_norm']),
        'normal_equation_defect': float(gm['normal_equation_defect']),
        'arnoldi_defect': float(arnoldi_defect),
        'hessenberg_defect': float(hessenberg_defect),
    }
    return result

Prompt

Calculate the deterministic sketched block CholeskyQR factorization of a tall full-rank matrix with a prescribed block partition and sketch. Let W∈Rn×mW\in\mathbb{R}^{n\times m} be partitioned into consecutive column blocks W(i)∈Rn×bW_{(i)}\in\mathbb{R}^{n\times b} for i=1,…,pi=1,\ldots,p, where m=pbm=pb. Let Θ∈Rk×n\Theta\in\mathbb{R}^{k\times n} with k≥mk\ge m. The admissible inputs are those for which every sketched residual Gram matrix in the following block recurrence is positive definite and the final full-space Gram matrix is positive definite. The returned pair Q,RQ,R is defined by the block recurrence over the given block order: the residual block for W(i)W_{(i)} is orthogonalized against all previously returned blocks using the sketched inner product induced by Θ\Theta; its diagonal triangular factor is the upper Cholesky factor with positive diagonal of the sketched residual Gram matrix; the block is normalized by that factor; and the completed factorization is followed by one full-space CholeskyQR cleanup that preserves W=QRW=QR and enforces full Euclidean orthonormality of QQ. The returned RR must be upper triangular with strictly positive diagonal. If any required positive-definiteness or shape condition fails, raise ValueError. Return a flat dictionary containing the cleaned factors, the cleaned sketch of QQ, and diagnostics. The relative QR residual is ∥W−QR∥F/max⁡(∥W∥F,ϵmach)\|W-QR\|_F/\max(\|W\|_F,\epsilon_{\rm mach}). The orthogonality defect is ∥Q⊤Q−I∥2\|Q^\top Q-I\|_2. The sketch orthogonality defect is ∥ΘQ⊤ΘQ−I∥2\|\Theta Q^\top\Theta Q-I\|_2 interpreted as ∥(ΘQ)⊤(ΘQ)−I∥2\| (\Theta Q)^\top(\Theta Q)-I\|_2. The triangular defect is the Frobenius norm of the strict lower triangle of RR divided by max⁡(∥R∥F,ϵmach)\max(\|R\|_F,\epsilon_{\rm mach}).

def rbgs_block_cholesky_qr(W, block_width, theta):
    '''
    Inputs:
        W : np.ndarray, shape (n, m), full-rank real matrix with m divisible by block_width
        block_width : int, width of each consecutive column block
        theta : np.ndarray, shape (k, n), real sketching matrix with k >= m
    Outputs:
        result : dict
            'Q' : np.ndarray, shape (n, m), cleaned orthonormal factor
            'R' : np.ndarray, shape (m, m), upper triangular factor with positive diagonal
            'thetaQ' : np.ndarray, shape (k, m), sketch of Q
            'relative_residual' : float, relative Frobenius defect of W = Q R
            'orthogonality_defect' : float, spectral norm of Q.T Q - I
            'sketch_orthogonality_defect' : float, spectral norm of (thetaQ).T thetaQ - I
            'triangular_defect' : float, relative strict-lower-triangular defect of R
            'min_diagonal_R' : float, minimum diagonal entry of R
            'rank_ok' : bool, whether all positive-definiteness checks succeeded
    '''
    return result

Background

Block Gram-Schmidt constructs a basis by repeatedly removing components in previously generated subspaces. In the sketched variant, the projection coefficients are governed by the lower-dimensional bilinear form ⟨x,y⟩Θ=(Θx)⊤(Θy)\langle x,y\rangle_\Theta=(\Theta x)^\top(\Theta y). The diagonal normalization of each residual block is a CholeskyQR condition: if VV is the residual block, the factor CC satisfies (ΘV)⊤(ΘV)=C⊤C(\Theta V)^\top(\Theta V)=C^\top C with CC upper triangular and positive diagonal, and the normalized block is VC−1VC^{-1}. Because sketched orthogonality and Euclidean orthogonality are not identical, the cleanup CholeskyQR step enforces the full-space invariant Q⊤Q≈IQ^\top Q\approx I while consistently updating RR so that the product QRQR continues to represent WW. Near-dependent blocks create small Cholesky pivots and expose rank or embedding failures.

Function signature

def rbgs_block_cholesky_qr(W, block_width, theta):
    return result

Published by the benchmark's authors. Opening these shows how the problem is solved and what the graders assert.

Test case 1 of 4
def test_case_1():
    import numpy as np
    import scipy.linalg as la
    # Tests W=QR, Q.T Q=I, positive-diagonal QR uniqueness from the block CholeskyQR equations; catches using sketched normalization without the required full cleanup.

    W = np.array([[2.0, -1.0, 0.5, 3.0], [1.5, 2.5, -2.0, 0.7], [-0.8, 1.2, 2.2, -1.1], [3.1, -0.4, 1.7, 2.6], [0.9, 2.8, -1.3, 1.4], [-2.2, 0.6, 3.4, -0.5], [1.1, -1.7, 0.8, 2.9]])
    theta = np.diag([1.2, 0.7, 1.5, 0.9, 1.1, 0.8, 1.3])
    theta[1, 0] = 0.23
    theta[3, 2] = -0.31
    theta[6, 4] = 0.17
    result = rbgs_block_cholesky_qr(W, 2, theta)
    q, r = la.qr(W, mode='economic')
    sgn = np.where(np.diag(r) < 0.0, -1.0, 1.0)
    q = q * sgn
    r = sgn[:, None] * r
    assert np.allclose(result['Q'], q, rtol=1e-9, atol=1e-10)
    assert np.allclose(result['R'], r, rtol=1e-9, atol=1e-10)
    assert abs(result['relative_residual']) < 1e-12
    expected = la.norm(
        result['Q'].T @ result['Q'] - np.eye(result['Q'].shape[1]), ord=2
    )
    assert np.isclose(
        result['orthogonality_defect'], expected, rtol=1e-12, atol=1e-15
    )
Test case 2 of 4
def test_case_2():
    import numpy as np
    import scipy.linalg as la
    # Uses a non-square distorted sketch in the sketched inner-product projection; catches implementations assuming theta is square/orthonormal or using theta.T theta incorrectly.

    W = np.array([[1.0, 2.0, -1.5, 0.7], [0.3, -1.2, 2.4, 3.1], [2.2, 0.4, -0.8, 1.9], [-1.7, 2.8, 0.6, -2.3], [3.5, -0.9, 1.1, 0.2], [0.8, 1.6, -2.6, 2.7], [-2.4, 0.5, 3.2, -1.4]])
    theta = np.array([[1.0, 0.2, -0.1, 0.0, 0.3, -0.2, 0.4], [0.1, 1.3, 0.2, -0.4, 0.0, 0.5, -0.1], [-0.3, 0.4, 1.1, 0.2, -0.5, 0.1, 0.2], [0.2, -0.1, 0.3, 1.4, 0.2, -0.3, 0.1], [0.5, 0.0, -0.2, 0.1, 1.2, 0.4, -0.3], [-0.2, 0.3, 0.1, -0.5, 0.2, 1.1, 0.6]])
    result = rbgs_block_cholesky_qr(W, 2, theta)
    q, r = la.qr(W, mode='economic')
    sgn = np.where(np.diag(r) < 0.0, -1.0, 1.0)
    q = q * sgn
    r = sgn[:, None] * r
    theta_q = theta @ result['Q']
    assert np.allclose(result['thetaQ'], theta_q, rtol=1e-12, atol=1e-12)
    expected_sketch_defect = la.norm(
        theta_q.T @ theta_q - np.eye(theta_q.shape[1]), ord=2
    )
    assert np.isclose(
        result['sketch_orthogonality_defect'],
        expected_sketch_defect,
        rtol=1e-12,
        atol=1e-12,
    )
    assert np.allclose(result['R'], r, rtol=1e-8, atol=1e-9)
    assert result['min_diagonal_R'] > 0.1
    assert abs(result['triangular_defect']) < 1e-14
Test case 3 of 4
def test_case_3():
    import numpy as np
    # Tests the block invariant W_(1:i)=Q_(1:i)R_(1:i,1:i) and strict upper triangular R for block_width=3; catches off-by-one block slicing.

    W = np.array([[1.2, -0.7, 2.1, 0.4, -1.5, 2.8], [2.3, 1.1, -0.6, 3.2, 0.9, -2.0], [-1.4, 2.7, 0.5, -0.8, 2.6, 1.3], [0.6, -2.2, 3.1, 1.7, -0.4, 0.8], [3.4, 0.3, -1.2, 2.5, 1.9, -0.7], [-2.1, 1.8, 0.9, -1.6, 3.0, 2.2], [0.5, 2.4, -2.8, 0.1, -1.1, 3.3], [1.7, -1.9, 2.6, -2.4, 0.7, 1.5]])
    theta = np.eye(8)
    theta[0, 3] = 0.37
    theta[4, 1] = -0.29
    theta[7, 5] = 0.41
    result = rbgs_block_cholesky_qr(W, 3, theta)
    Q = result['Q']; R = result['R']
    assert np.allclose(W, Q @ R, rtol=1e-10, atol=1e-10)
    assert np.allclose(Q[:, :3] @ R[:3, :3], W[:, :3], rtol=1e-10, atol=1e-10)
    assert np.allclose(Q.T @ Q, np.eye(6), rtol=1e-10, atol=1e-10)
    assert np.allclose(np.tril(R, -1), np.zeros((6, 6)), rtol=0.0, atol=1e-12)
    assert np.min(np.diag(R)) > 0.05
Test case 4 of 4
def test_case_4():
    import numpy as np
    import scipy.linalg as la
    # Discriminates the dangerous right-solve transpose error Q=V R^{-1} versus solving with R.T; the QR residual and unique positive R expose it.

    W = np.array([[4.0, 1.0, -2.0, 3.5], [1.2, -3.4, 2.2, 0.8], [-2.6, 2.9, 1.5, -1.1], [3.3, -0.7, 4.1, 2.4], [0.5, 3.8, -1.9, 1.7], [-1.8, 0.6, 2.7, -3.2], [2.1, -2.5, 0.4, 2.9]])
    theta = np.array([[1.4, 0.2, 0.0, -0.1, 0.3, 0.0, 0.1], [0.0, 1.1, -0.3, 0.2, 0.0, 0.4, -0.2], [0.2, -0.1, 1.6, 0.1, -0.4, 0.2, 0.0], [-0.3, 0.2, 0.1, 1.3, 0.5, -0.1, 0.2], [0.1, 0.0, -0.2, 0.3, 1.5, 0.2, -0.4], [0.2, -0.3, 0.4, 0.0, 0.1, 1.2, 0.3], [-0.1, 0.4, 0.0, -0.2, 0.3, 0.1, 1.7]])
    result = rbgs_block_cholesky_qr(W, 2, theta)
    q, r = la.qr(W, mode='economic')
    sgn = np.where(np.diag(r) < 0.0, -1.0, 1.0)
    r = sgn[:, None] * r
    assert np.allclose(result['R'], r, rtol=1e-9, atol=1e-10)
    assert np.allclose(W - result['Q'] @ result['R'], np.zeros_like(W), rtol=0.0, atol=1e-10)
    assert np.all(np.diag(result['R']) > 0.0)
    assert result['min_diagonal_R'] > 1.0
Solution
import numpy as np
import scipy.linalg as la


def _norm_fro(X):
    X = np.asarray(X)
    if X.size == 0:
        return 0.0
    return float(la.norm(X, ord='fro'))


def _norm_2(X):
    X = np.asarray(X)
    if X.size == 0:
        return 0.0
    return float(la.norm(X, ord=2))


def _chol_upper_spd(G, name):
    G = np.asarray(G, dtype=float)
    if G.ndim != 2 or G.shape[0] != G.shape[1]:
        raise ValueError(name + ' must be square')
    G = 0.5 * (G + G.T)
    if not np.all(np.isfinite(G)):
        raise ValueError(name + ' contains nonfinite entries')
    try:
        return la.cholesky(G, lower=False, check_finite=True)
    except la.LinAlgError as exc:
        raise ValueError(name + ' is not positive definite') from exc


def _right_solve_upper(V, R):
    if not np.all(np.isfinite(V)) or not np.all(np.isfinite(R)):
        raise ValueError('nonfinite triangular solve input')
    diagonal = np.abs(np.diag(R))
    scale = max(float(la.norm(R, ord=np.inf)), 1.0)
    tolerance = np.finfo(float).eps * scale

    if np.any(diagonal <= tolerance):
        raise ValueError('singular triangular factor')
    return la.solve_triangular(R, V.T, trans='T', lower=False, check_finite=True).T


def rbgs_block_cholesky_qr(W, block_width, theta):
    W = np.asarray(W, dtype=float)
    theta = np.asarray(theta, dtype=float)
    eps = np.finfo(float).eps
    if W.ndim != 2 or theta.ndim != 2:
        raise ValueError('W and theta must be two-dimensional')
    n, m = W.shape
    k, nt = theta.shape
    if nt != n:
        raise ValueError('theta has incompatible column dimension')
    if not isinstance(block_width, (int, np.integer)) or block_width <= 0:
        raise ValueError('block_width must be a positive integer')
    b = int(block_width)
    if m == 0 or m % b != 0:
        raise ValueError('number of columns of W must be positive and divisible by block_width')
    if k < m:
        raise ValueError('theta must have at least as many rows as W has columns')
    if not np.all(np.isfinite(W)) or not np.all(np.isfinite(theta)):
        raise ValueError('inputs contain nonfinite entries')
    p = m // b
    Q = np.zeros((n, m), dtype=float)
    R = np.zeros((m, m), dtype=float)
    thetaQ = np.zeros((k, m), dtype=float)
    for i in range(p):
        cs = slice(i * b, (i + 1) * b)
        V = W[:, cs].copy()
        TV = theta @ V
        for j in range(i):
            rs = slice(j * b, (j + 1) * b)
            C = thetaQ[:, rs].T @ TV
            R[rs, cs] = C
            V = V - Q[:, rs] @ C
            TV = TV - thetaQ[:, rs] @ C
        G = TV.T @ TV
        Rii = _chol_upper_spd(G, 'sketched block Gram matrix')
        R[cs, cs] = Rii
        Q[:, cs] = _right_solve_upper(V, Rii)
        thetaQ[:, cs] = _right_solve_upper(TV, Rii)
    Gfull = Q.T @ Q
    Cfull = _chol_upper_spd(Gfull, 'full cleanup Gram matrix')
    Q = _right_solve_upper(Q, Cfull)
    R = Cfull @ R
    thetaQ = theta @ Q
    rel = _norm_fro(W - Q @ R) / max(_norm_fro(W), eps)
    orth = _norm_2(Q.T @ Q - np.eye(m))
    sketch_orth = _norm_2(thetaQ.T @ thetaQ - np.eye(m))
    tri = _norm_fro(np.tril(R, -1)) / max(_norm_fro(R), eps)
    mind = float(np.min(np.diag(R))) if m else 0.0
    result = {
        'Q': Q,
        'R': R,
        'thetaQ': thetaQ,
        'relative_residual': float(rel),
        'orthogonality_defect': float(orth),
        'sketch_orthogonality_defect': float(sketch_orth),
        'triangular_defect': float(tri),
        'min_diagonal_R': mind,
        'rank_ok': bool(mind > 0.0)
    }
    return result

Prompt

Calculate the sketched Petrov-Galerkin minimum-residual correction for a supplied trial basis on the stiff Chebyshev-Lobatto collocation operator. Let n≥4n\ge 4 and let xj=cos⁡(πj/(n−1))x_j=\cos(\pi j/(n-1)) for j=0,…,n−1j=0,\ldots,n-1. Let DD be the nodal polynomial differentiation matrix on these nodes and D(2)=DDD^{(2)}=DD. For ν>0\nu>0, a∈Ra\in\mathbb{R}, and ρ≥0\rho\ge 0, define A∈Rn×nA\in\mathbb{R}^{n\times n} by the interior operator −νD(2)+aD+ρI-\nu D^{(2)}+aD+\rho I on rows 1,…,n−21,\ldots,n-2 and Dirichlet boundary rows A0,:=e0⊤A_{0,:}=e_0^\top, An−1,:=en−1⊤A_{n-1,:}=e_{n-1}^\top. Given B,X0∈Rn×sB,X_0\in\mathbb{R}^{n\times s}, a trial basis Q∈Rn×mQ\in\mathbb{R}^{n\times m}, and a sketch Θ∈Rk×n\Theta\in\mathbb{R}^{k\times n}, the returned coefficient matrix Y∈Rm×sY\in\mathbb{R}^{m\times s} is the minimum-Frobenius-norm minimizer of min⁡Y∈Rm×s∥Θ(A(X0+QY)−B)∥F2. \min_{Y\in\mathbb{R}^{m\times s}}\|\Theta(A(X_0+QY)-B)\|_F^2. Let Z=ΘAQZ=\Theta A Q. Singular values satisfying σi(Z)≤εmachσmax⁡(Z)\sigma_i(Z)\leq\varepsilon_{\mathrm{mach}}\sigma_{\max}(Z) are treated as zero under the Moore-Penrose minimum-Frobenius-norm convention. Both numerically full-rank and rank-deficient cases are admissible. Shape mismatches, nonfinite inputs, n<4n<4, ν≤0\nu\le 0, or ρ<0\rho<0 must raise ValueError. Let rΘ=Θ(AX−B)r_\Theta=\Theta(AX-B) and Z=ΘAQZ=\Theta A Q. Define \texttt{normal_equation_defect} = \frac{\|Z^\top r_\Theta\|_F} {\max(\|Z\|_F\|r_\Theta\|_F,\varepsilon_{\mathrm{mach}})} and \texttt{operator_condition} = \frac{\sigma_{\max}(A)} {\max(\sigma_{\min}(A),\varepsilon_{\mathrm{mach}})}. Return a flat dictionary with the approximation, coefficient matrix, residual diagnostics, and these two quantities.

def chebyshev_sketched_trial_gmres(n, B, X0, Q, theta, nu, adv, rho):
    '''
    Inputs:
        n : int, number of Chebyshev-Lobatto nodes, n >= 4
        B : np.ndarray, shape (n, s), right-hand side block
        X0 : np.ndarray, shape (n, s), initial guess block
        Q : np.ndarray, shape (n, m), trial basis
        theta : np.ndarray, shape (k, n), real sketching matrix
        nu : float, positive diffusion coefficient
        adv : float, advection coefficient
        rho : float, nonnegative reaction coefficient
    Outputs:
        result : dict
            'X' : np.ndarray, shape (n, s), sketched minimum-residual approximation
            'Y' : np.ndarray, shape (m, s), minimum-Frobenius-norm coefficient matrix
            'sketched_residual_norm' : float, Frobenius norm of theta @ (A @ X - B)
            'true_residual_norm' : float, Frobenius norm of A @ X - B
            'normal_equation_defect' : float, relative norm of the sketched Petrov-Galerkin defect
            'operator_condition' : float, guarded spectral condition number of A
    '''
    return result

Background

A sketched minimum-residual projection replaces the full residual norm by a lower-dimensional residual norm induced by Θ\Theta. The trial solution has the constrained form X=X0+QYX=X_0+QY, so the unknown is the coefficient matrix YY. The first-order optimality condition in sketched space is the Petrov-Galerkin relation (ΘAQ)⊤Θ(AX−B)=0(\Theta A Q)^\top\Theta(AX-B)=0 whenever the least-squares problem is differentiable in the usual Euclidean sense. The Chebyshev-Lobatto grid is strongly nonuniform, with nodes clustered near ±1\pm 1; as a result, the differentiation matrices produce a nonnormal and stiff discretized operator. The diffusion term −νD(2)-\nu D^{(2)} and the advection term aDaD interact through this nonuniform collocation geometry, making the residual projection sensitive to both basis conditioning and sketch distortion.

Function signature

def chebyshev_sketched_trial_gmres(n, B, X0, Q, theta, nu, adv, rho):
    return result

Published by the benchmark's authors. Opening these shows how the problem is solved and what the graders assert.

Test case 1 of 4
def test_case_1():
    import numpy as np
    import scipy.linalg as la
    # Tests the sketched Petrov-Galerkin least-squares equation (Theta A Q)Y=Theta(B-AX0); catches sign errors in B-AX0.

    def Aop(n, nu, adv, rho):
        j = np.arange(n, dtype=float); x = np.cos(np.pi * j / float(n - 1)); c = np.ones(n); c[0] = 2.0; c[-1] = 2.0; c = c * ((-1.0) ** np.arange(n)); dX = x[:, None] - x[None, :]; D = np.zeros((n, n)); mask = ~np.eye(n, dtype=bool); C = c[:, None] / c[None, :]; D[mask] = C[mask] / dX[mask]; D[np.diag_indices(n)] = -np.sum(D, axis=1); A = -nu * (D @ D) + adv * D + rho * np.eye(n); A[0, :] = 0.0; A[0, 0] = 1.0; A[-1, :] = 0.0; A[-1, -1] = 1.0; return A
    n = 7; nu = 0.08; adv = 1.7; rho = 0.9
    B = np.array([[1.0, -0.5], [2.2, 1.3], [-1.1, 2.4], [0.7, -2.0], [3.1, 0.6], [-0.8, 1.9], [1.5, -1.4]])
    X0 = np.array([[0.2, -0.1], [-0.3, 0.5], [0.4, 0.7], [-0.6, 0.2], [0.8, -0.4], [0.1, 0.9], [-0.5, 0.3]])
    Q = np.array([[1.0, 0.2, -0.3], [0.4, 1.1, 0.5], [-0.7, 0.3, 1.2], [1.5, -0.8, 0.6], [0.2, 1.4, -1.0], [-1.1, 0.7, 0.9], [0.6, -0.4, 1.3]])
    theta = np.eye(n); theta[2, 0] = 0.31; theta[5, 3] = -0.27
    result = chebyshev_sketched_trial_gmres(n, B, X0, Q, theta, nu, adv, rho)
    A = Aop(n, nu, adv, rho); Z = theta @ (A @ Q); S = theta @ (B - A @ X0)
    Yexp, _, _, _ = la.lstsq(Z, S, lapack_driver='gelsd')
    Xexp = X0 + Q @ Yexp
    assert np.allclose(result['Y'], Yexp, rtol=1e-10, atol=1e-10)
    assert np.allclose(result['X'], Xexp, rtol=1e-10, atol=1e-10)
    assert result['normal_equation_defect'] < 1e-12
    svals = la.svdvals(A)
    expected_condition = svals[0] / max(svals[-1], np.finfo(float).eps)
    assert np.isclose(
        result['operator_condition'], expected_condition,
        rtol=1e-12, atol=1e-12
    )
Test case 2 of 4
def test_case_2():
    import numpy as np
    import scipy.linalg as la
    # Exercises Chebyshev-Lobatto nonuniform differentiation with active diffusion/advection/reaction; catches replacing the clustered cosine grid by a uniform grid.

    def Aop(n, nu, adv, rho):
        j = np.arange(n, dtype=float); x = np.cos(np.pi * j / float(n - 1)); c = np.ones(n); c[0] = 2.0; c[-1] = 2.0; c = c * ((-1.0) ** np.arange(n)); dX = x[:, None] - x[None, :]; D = np.zeros((n, n)); mask = ~np.eye(n, dtype=bool); C = c[:, None] / c[None, :]; D[mask] = C[mask] / dX[mask]; D[np.diag_indices(n)] = -np.sum(D, axis=1); A = -nu * (D @ D) + adv * D + rho * np.eye(n); A[0, :] = 0.0; A[0, 0] = 1.0; A[-1, :] = 0.0; A[-1, -1] = 1.0; return A
    n = 8; nu = 0.05; adv = -2.4; rho = 1.2
    B = np.array([[2.0], [-1.5], [3.2], [0.4], [-2.1], [1.7], [0.8], [-0.6]])
    X0 = np.array([[0.1], [0.3], [-0.2], [0.5], [-0.7], [0.9], [-0.4], [0.2]])
    Q = np.array([[1.0, -0.2, 0.3], [0.5, 1.4, -0.6], [-1.2, 0.7, 1.1], [0.8, -1.5, 0.4], [1.6, 0.2, -1.3], [-0.9, 1.0, 0.8], [0.3, -0.7, 1.5], [1.1, 0.6, -0.4]])
    theta = np.array([[1.0, 0.2, 0.0, -0.1, 0.3, 0.0, 0.1, -0.2], [0.0, 1.3, -0.2, 0.2, 0.0, 0.4, -0.1, 0.3], [0.2, -0.1, 1.1, 0.0, -0.3, 0.2, 0.0, 0.1], [-0.1, 0.3, 0.2, 1.4, 0.1, -0.2, 0.4, 0.0], [0.3, 0.0, -0.4, 0.1, 1.2, 0.2, -0.2, 0.5], [0.0, 0.2, 0.1, -0.3, 0.4, 1.5, 0.3, -0.1]])
    result = chebyshev_sketched_trial_gmres(n, B, X0, Q, theta, nu, adv, rho)
    A = Aop(n, nu, adv, rho); Z = theta @ (A @ Q); S = theta @ (B - A @ X0); Yexp, _, _, _ = la.lstsq(Z, S, lapack_driver='gelsd')
    assert np.allclose(result['Y'], Yexp, rtol=1e-10, atol=1e-10)
    expected_norm = la.norm(
        theta @ (A @ (X0 + Q @ Yexp) - B),
        ord='fro',
    )

    assert np.isclose(
        result['sketched_residual_norm'],
        expected_norm,
        rtol=1e-10,
        atol=1e-9,
    )
    assert result['true_residual_norm'] > 1e-3
    try:
        chebyshev_sketched_trial_gmres(
            n, B, X0, Q, theta, nu, adv, -0.3
        )
    except ValueError:
        pass
    else:
        raise AssertionError('rho < 0 must raise ValueError')
Test case 3 of 4
def test_case_3():
    import numpy as np
    import scipy.linalg as la
    # Verifies the specific normal-equation invariant (Theta A Q).T Theta(A X-B)=0; catches using the unsketched residual objective.

    def Aop(n, nu, adv, rho):
        j = np.arange(n, dtype=float); x = np.cos(np.pi * j / float(n - 1)); c = np.ones(n); c[0] = 2.0; c[-1] = 2.0; c = c * ((-1.0) ** np.arange(n)); dX = x[:, None] - x[None, :]; D = np.zeros((n, n)); mask = ~np.eye(n, dtype=bool); C = c[:, None] / c[None, :]; D[mask] = C[mask] / dX[mask]; D[np.diag_indices(n)] = -np.sum(D, axis=1); A = -nu * (D @ D) + adv * D + rho * np.eye(n); A[0, :] = 0.0; A[0, 0] = 1.0; A[-1, :] = 0.0; A[-1, -1] = 1.0; return A
    n = 9; nu = 0.11; adv = 2.2; rho = 0.4
    B = np.reshape(np.linspace(-2.0, 2.5, 18), (9, 2))
    X0 = np.reshape(np.linspace(0.7, -1.1, 18), (9, 2))
    Q = np.reshape(np.sin(np.linspace(0.2, 5.4, 36)), (9, 4)) + 0.3 * np.reshape(np.cos(np.linspace(0.1, 3.7, 36)), (9, 4))
    theta = np.vstack([np.eye(9)[:7], np.array([[0.2, -0.3, 0.4, 0.1, -0.2, 0.5, 0.0, 0.3, -0.1]])])
    result = chebyshev_sketched_trial_gmres(n, B, X0, Q, theta, nu, adv, rho)
    A = Aop(n, nu, adv, rho); Z = theta @ (A @ Q); sk_res = theta @ (A @ result['X'] - B); normal = Z.T @ sk_res
    denom = max(la.norm(Z, ord='fro') * la.norm(sk_res, ord='fro'), np.finfo(float).eps)
    assert abs(result['normal_equation_defect'] - la.norm(normal, ord='fro') / denom) < 1e-12
    assert result['sketched_residual_norm'] > 1e-6
Test case 4 of 4
def test_case_4():
    import numpy as np
    import scipy.linalg as la
    # Tests Moore-Penrose minimum-Frobenius-norm convention for rank-deficient Theta A Q; catches solvers using normal equations without pseudoinverse handling.

    def Aop(n, nu, adv, rho):
        j = np.arange(n, dtype=float); x = np.cos(np.pi * j / float(n - 1)); c = np.ones(n); c[0] = 2.0; c[-1] = 2.0; c = c * ((-1.0) ** np.arange(n)); dX = x[:, None] - x[None, :]; D = np.zeros((n, n)); mask = ~np.eye(n, dtype=bool); C = c[:, None] / c[None, :]; D[mask] = C[mask] / dX[mask]; D[np.diag_indices(n)] = -np.sum(D, axis=1); A = -nu * (D @ D) + adv * D + rho * np.eye(n); A[0, :] = 0.0; A[0, 0] = 1.0; A[-1, :] = 0.0; A[-1, -1] = 1.0; return A
    n = 7; nu = 0.07; adv = 1.3; rho = 0.6
    B = np.array([[1.4], [-0.8], [2.6], [0.9], [-1.7], [3.0], [0.2]])
    X0 = np.array([[-0.2], [0.5], [0.1], [-0.6], [0.8], [-0.4], [0.3]])
    q1 = np.array([1.0, 0.5, -0.7, 1.2, -1.1, 0.3, 0.9])[:, None]
    q2 = np.array([-0.4, 1.3, 0.6, -0.8, 0.2, 1.1, -1.5])[:, None]
    Q = np.hstack([q1, q2, q1 + 2.0 * q2])
    theta = np.eye(n)
    result = chebyshev_sketched_trial_gmres(n, B, X0, Q, theta, nu, adv, rho)
    A = Aop(n, nu, adv, rho); Z = theta @ (A @ Q); S = theta @ (B - A @ X0)
    Yexp, _, _, _ = la.lstsq(
        Z,
        S,
        cond=np.finfo(float).eps,
        lapack_driver='gelsd',
    )
    assert np.allclose(result['Y'], Yexp, rtol=1e-10, atol=1e-10)
    assert np.allclose(result['X'], X0 + Q @ Yexp, rtol=1e-10, atol=1e-10)
    assert result['normal_equation_defect'] < 1e-11
    assert la.norm(result['Y'], ord='fro') > 1e-4
Solution
import numpy as np
import scipy.linalg as la


def _norm_fro(X):
    X = np.asarray(X)
    if X.size == 0:
        return 0.0
    return float(la.norm(X, ord='fro'))


def _chebyshev_lobatto_operator(n, nu, adv, rho):
    if not isinstance(n, (int, np.integer)) or int(n) < 4:
        raise ValueError('n must be an integer at least 4')
    n = int(n)
    nu = float(nu)
    adv = float(adv)
    rho = float(rho)
    if not np.isfinite(nu) or not np.isfinite(adv) or not np.isfinite(rho):
        raise ValueError('operator parameters must be finite')
    if nu <= 0.0 or rho < 0.0:
        raise ValueError('nu must be positive and rho must be nonnegative')
    j = np.arange(n, dtype=float)
    x = np.cos(np.pi * j / float(n - 1))
    c = np.ones(n, dtype=float)
    c[0] = 2.0
    c[-1] = 2.0
    c = c * ((-1.0) ** np.arange(n))
    X = x[:, None]
    dX = X - X.T
    D = np.zeros((n, n), dtype=float)
    mask = ~np.eye(n, dtype=bool)
    C = c[:, None] / c[None, :]
    D[mask] = C[mask] / dX[mask]
    D[np.diag_indices(n)] = -np.sum(D, axis=1)
    D2 = D @ D
    A = -nu * D2 + adv * D + rho * np.eye(n)
    A[0, :] = 0.0
    A[0, 0] = 1.0
    A[-1, :] = 0.0
    A[-1, -1] = 1.0
    if not np.all(np.isfinite(A)):
        raise ValueError('constructed operator is nonfinite')
    return A


def chebyshev_sketched_trial_gmres(n, B, X0, Q, theta, nu, adv, rho):
    B = np.asarray(B, dtype=float)
    X0 = np.asarray(X0, dtype=float)
    Q = np.asarray(Q, dtype=float)
    theta = np.asarray(theta, dtype=float)
    eps = np.finfo(float).eps
    if B.ndim != 2 or X0.ndim != 2 or Q.ndim != 2 or theta.ndim != 2:
        raise ValueError('B, X0, Q, and theta must be two-dimensional')
    if B.shape != X0.shape:
        raise ValueError('B and X0 must have the same shape')
    if B.shape[0] != int(n) or Q.shape[0] != int(n) or theta.shape[1] != int(n):
        raise ValueError('incompatible n dimension')
    if not np.all(np.isfinite(B)) or not np.all(np.isfinite(X0)) or not np.all(np.isfinite(Q)) or not np.all(np.isfinite(theta)):
        raise ValueError('inputs contain nonfinite entries')
    A = _chebyshev_lobatto_operator(int(n), nu, adv, rho)
    Z = theta @ (A @ Q)
    S = theta @ (B - A @ X0)
    if not np.all(np.isfinite(Z)) or not np.all(np.isfinite(S)):
        raise ValueError('least-squares matrix or right-hand side is nonfinite')
    Y, _, _, _ = la.lstsq(
        Z, S, cond=eps, lapack_driver='gelsd', check_finite=True
    )
    X = X0 + Q @ Y
    residual = A @ X - B
    sk_res = theta @ residual
    normal = Z.T @ sk_res
    normal_denom = max(_norm_fro(Z) * _norm_fro(sk_res), eps)
    svals = la.svdvals(A)
    if svals.size == 0 or not np.all(np.isfinite(svals)):
        raise ValueError('operator singular values are invalid')
    cond = float(np.max(svals) / max(np.min(svals), eps))
    result = {
        'X': X,
        'Y': Y,
        'sketched_residual_norm': float(_norm_fro(sk_res)),
        'true_residual_norm': float(_norm_fro(residual)),
        'normal_equation_defect': float(_norm_fro(normal) / normal_denom),
        'operator_condition': cond
    }
    return result