Constrained Least Squares and Advanced Tridiagonal

Demonstrates non-negative, equality-constrained, and bounded least-squares solvers, the cyclic and block tridiagonal solvers, complex general linear system solve, and complex triangular solve.

Functions covered: linearsystems.lsqnr(), linearsystems.lsqcr(), linearsystems.lsbrr(), linearsystems.cyclic_tridiag_solve(), linearsystems.block_tridiag_solve(), linearsystems.lscrg(), linearsystems.lsltc().

Example Code

"""Constrained Least Squares and Advanced Tridiagonal example.

Demonstrates:
- lsqnr: non-negative least squares (fit non-negative spectral amplitudes)
- lsqcr: constrained least squares with equality constraints (polynomial fit
  constrained to pass through two specified points)
- lsbrr: bounded least squares (polynomial fit with coefficient bounds)
- cyclic_tridiag_solve: periodic 1D heat equation (cyclic boundary)
- block_tridiag_solve: 2-block tridiagonal system
- lscrg: solve a complex general linear system
- lsltc: solve a complex triangular system

Outputs:
- Constraint satisfaction and residuals printed to stdout
- SVG: LS solutions comparison + cyclic tridiagonal solution
  saved to test_output/example_imsl_constrained_ls.svg
"""

from __future__ import annotations

from pathlib import Path
from typing import Dict

import matplotlib.pyplot as plt
import numpy as np

from linearsystems import (
    lsqnr,
    lsqcr,
    lsbrr,
    cyclic_tridiag_solve,
    block_tridiag_solve,
    lscrg,
    lsltc,
)


def run_demo_imsl_constrained_ls() -> Dict[str, object]:
    """Run constrained least squares and advanced tridiagonal demo.

    Builds non-negative, equality-constrained, and bounded least-squares
    problems plus a cyclic and block tridiagonal system.  Also demonstrates
    complex linear system solve (lscrg) and complex triangular solve (lsltc).

    Args:
        None

    Returns:
        Dict[str, object]: Keys: ``nnls_res``, ``constrained_res``,
            ``bounded_res``, ``cyclic_max_err``, ``complex_residual``,
            ``plot_path`` (str).
    """
    rng = np.random.default_rng(99)

    # ------------------------------------------------------------------
    # 1. Non-negative LS  (lsqnr)
    # ------------------------------------------------------------------
    # Recover non-negative spectral amplitudes from noisy observations.
    # True signal: 3 non-negative components [2.5, 0.0, 1.8]
    n_features = 5
    m_obs = 20
    A_nn = rng.standard_normal((m_obs, n_features))
    x_nn_true = np.array([2.5, 0.0, 1.8, 0.0, 0.6])
    b_nn = A_nn @ x_nn_true + 0.05 * rng.standard_normal(m_obs)

    x_nnls, nnls_norm = lsqnr(A_nn, b_nn)
    nnls_res = float(np.linalg.norm(A_nn @ x_nnls - b_nn))

    print("\nConstrained Least Squares and Advanced Tridiagonal")
    print("=" * 65)
    print(f"\n1. Non-negative LS (lsqnr)  — {m_obs}×{n_features} system")
    print(f"   {'Component':<12} {'x_true':>10} {'x_nnls':>10}")
    for i in range(n_features):
        print(f"   x[{i}]        {x_nn_true[i]:>10.4f} {x_nnls[i]:>10.4f}")
    print(f"   ‖Ax − b‖₂ = {nnls_res:.4e}  (all x ≥ 0: {np.all(x_nnls >= -1e-12)})")

    # ------------------------------------------------------------------
    # 2. Constrained LS  (lsqcr)
    # ------------------------------------------------------------------
    # Polynomial fit (degree 3) with constraint that p(0)=1 and p(1)=2.
    t = np.linspace(0, 1, 30)
    degree = 3
    A_cr = np.column_stack([t ** k for k in range(degree + 1)])  # (30, 4)
    y_true = 1.0 + 3.0 * t - 2.0 * t ** 2 + 0.5 * t ** 3
    b_cr = y_true + 0.02 * rng.standard_normal(len(t))

    # Constraints: p(0) = [1,0,0,0]@x = 1;  p(1) = [1,1,1,1]@x = 2
    C = np.array([[1, 0, 0, 0], [1, 1, 1, 1]], dtype=float)
    d = np.array([1.0, 2.0])
    x_cr, _, _, _ = lsqcr(A_cr, b_cr, constraints={"C": C, "d": d})
    constrained_res = float(np.linalg.norm(A_cr @ x_cr - b_cr))
    constraint_err = float(np.linalg.norm(C @ x_cr - d))

    print(f"\n2. Constrained LS (lsqcr)  — polynomial fit degree {degree}")
    print(f"   Coefficients: {x_cr}")
    print(f"   ‖Ax − b‖₂         = {constrained_res:.4e}")
    print(f"   Constraint error  = {constraint_err:.4e}  (should be ≈0)")
    print(f"   p(0) = {C[0] @ x_cr:.4f}  (target 1.0)")
    print(f"   p(1) = {C[1] @ x_cr:.4f}  (target 2.0)")

    # ------------------------------------------------------------------
    # 3. Bounded LS  (lsbrr)
    # ------------------------------------------------------------------
    # Fit a cubic with coefficients in [-1, 3].
    lb = np.full(degree + 1, -1.0)
    ub = np.full(degree + 1, 3.0)
    x_br, br_norm = lsbrr(A_cr, b_cr, lb=lb, ub=ub)
    bounded_res = float(np.linalg.norm(A_cr @ x_br - b_cr))
    bounds_ok = np.all(x_br >= lb - 1e-6) and np.all(x_br <= ub + 1e-6)

    print(f"\n3. Bounded LS (lsbrr)  — coefficients in [-1, 3]")
    print(f"   Coefficients: {x_br}")
    print(f"   ‖Ax − b‖₂     = {bounded_res:.4e}")
    print(f"   Bounds satisfied: {bounds_ok}")

    # ------------------------------------------------------------------
    # 4. Cyclic tridiagonal solve
    # ------------------------------------------------------------------
    # Periodic heat equation: (alpha*I - Laplacian) u = f on a cycle of n=12.
    # Using alpha=4 ensures strict diagonal dominance and non-singularity.
    # Eigenvalues = 4 - 2*cos(2*pi*k/n) >= 2 > 0 for all k.
    nc = 12
    alpha = 4.0
    main_c = alpha * np.ones(nc)
    off_c = -1.0 * np.ones(nc)       # length n for cyclic (wrap-around)
    b_cyclic = np.zeros(nc)
    b_cyclic[nc // 2] = 1.0           # unit source at middle node

    u_cyclic = cyclic_tridiag_solve(off_c, main_c, off_c, b_cyclic)

    # Verify by building explicit cyclic matrix and checking residual
    A_cyc = np.diag(main_c) + np.diag(off_c[1:], -1) + np.diag(off_c[:-1], 1)
    A_cyc[0, -1] = off_c[0]
    A_cyc[-1, 0] = off_c[-1]
    cyclic_res = float(np.linalg.norm(A_cyc @ u_cyclic - b_cyclic))
    cyclic_max_err = float(np.max(np.abs(A_cyc @ u_cyclic - b_cyclic)))

    print(f"\n4. cyclic_tridiag_solve  (n={nc} periodic nodes, alpha={alpha})")
    print(f"   ‖Au − b‖₂    = {cyclic_res:.4e}")
    print(f"   max |error|  = {cyclic_max_err:.4e}")

    # ------------------------------------------------------------------
    # 5. Block tridiagonal solve
    # ------------------------------------------------------------------
    # 2 blocks of size 2: D0, U0 | L1, D1
    D0 = np.array([[4.0, 1.0], [1.0, 3.0]])
    U0 = np.array([[0.5, 0.0], [0.0, 0.5]])
    L1 = np.array([[0.5, 0.0], [0.0, 0.5]])
    D1 = np.array([[4.0, -1.0], [-1.0, 3.0]])
    blocks = [(None, D0, U0), (L1, D1, None)]
    b_bt = np.array([1.0, 2.0, 3.0, 4.0])
    x_bt = block_tridiag_solve(blocks, b_bt)

    # Verify
    A_bt_full = np.block([[D0, U0], [L1, D1]])
    bt_res = float(np.linalg.norm(A_bt_full @ x_bt - b_bt))
    print(f"\n5. block_tridiag_solve  (2 blocks × 2×2)")
    print(f"   x = {x_bt}")
    print(f"   ‖Ax − b‖₂ = {bt_res:.4e}")

    # ------------------------------------------------------------------
    # 6. Complex linear system  (lscrg)
    # ------------------------------------------------------------------
    nc_dim = 4
    Ac = rng.standard_normal((nc_dim, nc_dim)) + 1j * rng.standard_normal((nc_dim, nc_dim))
    Ac = Ac @ Ac.conj().T + nc_dim * np.eye(nc_dim)   # Hermitian positive definite
    xc_true = rng.standard_normal(nc_dim) + 1j * rng.standard_normal(nc_dim)
    bc = Ac @ xc_true
    result_crg = lscrg(Ac, bc)
    complex_residual = result_crg.residual

    print(f"\n6. lscrg (complex general system, n={nc_dim})")
    print(f"   ‖Ax − b‖₂ = {complex_residual:.4e}")

    # ------------------------------------------------------------------
    # 7. Complex triangular solve  (lsltc) — upper triangular
    # ------------------------------------------------------------------
    Uc = np.triu(rng.standard_normal((nc_dim, nc_dim)) + 1j * rng.standard_normal((nc_dim, nc_dim)))
    np.fill_diagonal(Uc, np.abs(np.diag(Uc)) + 2.0)
    bc2 = Uc @ xc_true
    result_ltc = lsltc(Uc, bc2, lower=False)
    ltc_residual = result_ltc.residual

    print(f"\n7. lsltc (complex upper triangular, n={nc_dim})")
    print(f"   ‖Ux − b‖₂ = {ltc_residual:.4e}")

    # ------------------------------------------------------------------
    # Plots
    # ------------------------------------------------------------------
    output_dir = Path("test_output")
    output_dir.mkdir(parents=True, exist_ok=True)
    plot_path = output_dir / "example_imsl_constrained_ls.svg"

    fig, axes = plt.subplots(1, 3, figsize=(18, 6))

    # (0) LS solutions comparison on polynomial fit
    t_fine = np.linspace(0, 1, 200)
    A_fine = np.column_stack([t_fine ** k for k in range(degree + 1)])
    y_cr_fit = A_fine @ x_cr
    y_br_fit = A_fine @ x_br
    ax = axes[0]
    ax.scatter(t, b_cr, s=12, color="grey", alpha=0.5, label="noisy data")
    ax.plot(t_fine, A_fine @ np.array([1, 3, -2, 0.5]), "k-", linewidth=2,
            label="true poly")
    ax.plot(t_fine, y_cr_fit, "b--", linewidth=2, label="lsqcr (constrained)")
    ax.plot(t_fine, y_br_fit, "r-.", linewidth=2, label="lsbrr (bounded)")
    ax.axvline(0, color="b", alpha=0.3, linewidth=0.8)
    ax.axvline(1, color="b", alpha=0.3, linewidth=0.8)
    ax.set_xlabel("t")
    ax.set_ylabel("y")
    ax.set_title("Polynomial fit: constrained vs bounded LS")
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)

    # (1) Cyclic tridiagonal solution
    ax2 = axes[1]
    theta = np.linspace(0, 2 * np.pi, nc, endpoint=False)
    ax2.plot(theta, u_cyclic, "b-o", linewidth=2, ms=6, label="cyclic solution")
    ax2.axvline(theta[nc // 2], color="r", linestyle="--", alpha=0.7, label="source node")
    ax2.set_xlabel("Node angle θ (radians)")
    ax2.set_ylabel("u(θ)")
    ax2.set_title(f"Cyclic tridiagonal: (4I - L)u = f  (n={nc})")
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    # (2) NNLS result bar chart
    ax3 = axes[2]
    idx = np.arange(n_features)
    ax3.bar(idx - 0.2, x_nn_true, width=0.35, label="x_true", color="#1f77b4", alpha=0.85)
    ax3.bar(idx + 0.2, x_nnls, width=0.35, label="x_nnls", color="#ff7f0e", alpha=0.85)
    ax3.set_xticks(idx)
    ax3.set_xticklabels([f"x[{i}]" for i in idx])
    ax3.set_title("Non-negative LS (lsqnr) recovery")
    ax3.set_ylabel("Coefficient value")
    ax3.legend()
    ax3.axhline(0, color="k", linewidth=0.8)
    ax3.grid(True, axis="y", alpha=0.3)

    fig.suptitle("Constrained Least Squares and Advanced Tridiagonal", fontsize=13)
    fig.tight_layout()
    fig.savefig(plot_path, format="svg")
    plt.close(fig)
    print(f"\nPlot saved to: {plot_path}")

    return {
        "nnls_res": nnls_res,
        "constrained_res": constrained_res,
        "bounded_res": bounded_res,
        "cyclic_max_err": cyclic_max_err,
        "complex_residual": complex_residual,
        "plot_path": str(plot_path),
    }


if __name__ == "__main__":
    run_demo_imsl_constrained_ls()

Plot Output

Constrained least-squares and advanced tridiagonal results

Left: polynomial fits — constrained (lsqcr) vs bounded (lsbrr). Centre: cyclic tridiagonal solution on 12-node periodic domain. Right: non-negative least-squares (lsqnr) coefficient recovery.

Console Output

Constrained Least Squares and Advanced Tridiagonal
=================================================================

1. Non-negative LS (lsqnr)  — 20×5 system
   Component        x_true     x_nnls
   x[0]            2.5000     2.4999
   x[1]            0.0000     0.0038
   x[2]            1.8000     1.7573
   x[3]            0.0000     0.0116
   x[4]            0.6000     0.6186
   ‖Ax − b‖₂ = 1.9360e-01  (all x ≥ 0: True)

2. Constrained LS (lsqcr)  ΓÇö polynomial fit degree 3
   Coefficients: [ 1.          1.94475114  2.19543491 -3.14018605]
   ‖Ax − b‖₂         = 7.7036e-01
   Constraint error  = 2.2204e-16  (should be Γëê0)
   p(0) = 1.0000  (target 1.0)
   p(1) = 2.0000  (target 2.0)

3. Bounded LS (lsbrr)  ΓÇö coefficients in [-1, 3]
   Coefficients: [ 1.028139    2.61999283 -1.         -0.17710188]
   ‖Ax − b‖₂     = 1.0922e-01
   Bounds satisfied: True

4. cyclic_tridiag_solve  (n=12 periodic nodes, alpha=4.0)
   ‖Au − b‖₂    = 1.3743e-16
   max |error|  = 1.1102e-16

5. block_tridiag_solve  (2 blocks × 2×2)
   x = [0.00758808 0.38807588 1.16314363 1.65636856]
   ‖Ax − b‖₂ = 9.9301e-16

6. lscrg (complex general system, n=4)
   ‖Ax − b‖₂ = 1.9860e-15

7. lsltc (complex upper triangular, n=4)
   ‖Ux − b‖₂ = 4.4409e-16

Plot saved to: test_output\example_imsl_constrained_ls.svg