Chapter 9 Experiment 9: Visualizing the Steps of Gram-Schmidt Orthogonalization

import numpy as np
import plotly.graph_objects as go
import ipywidgets as widgets
from IPython.display import display

np.set_printoptions(precision=4, suppress=True, linewidth=100)

# --- Color scheme: book palette ---
C_BG   = "#F8F8F8"
C_GRID = "#D6D6D6"
C_AXIS = "#000000"
C_V1   = "#57068C"   # violet: input vector v0
C_V2   = "#006385"   # deep blue: input vector v1
C_V3   = "#8900E1"   # Ultra Violet: input vector v2
C_T1   = "#2AD2C9"   # teal: orthonormal basis vector e0
C_T2   = "#FF5D47"   # orange: orthonormal basis vector e1 / projection
C_T3   = "#57068C"   # orthonormal basis vector e2
C_AUX  = "#AB82C5"   # light violet: degenerate residual

VEC_INPUT_COLORS = [C_V1, C_V2, C_V3]   # input vectors (dashed)
VEC_BASIS_COLORS = [C_T1, C_T2, C_T3]   # orthonormal basis (solid)
PROJ_COLOR       = C_T2                  # projection arrows
RESIDUE_COLOR    = C_V1                  # residual arrows

# ============================================================
# Core algorithm
# ============================================================

def gram_schmidt_steps(vecs):
    """
    Gram-Schmidt orthogonalization with a record of the steps.

    Returns
    -------
    basis : list[ndarray]  orthonormal basis vectors
    steps : list[dict]
        each step contains:
          v          the original input vector
          w          the residual (after subtracting all projections; not normalized)
          u          the normalized basis vector (the zero vector if degenerate)
          projs      list[{'j', 'coeff', 'vec'}]  the projection components
          norm_w     norm of the residual
          degenerate whether the step is degenerate (nearly linearly dependent)
    """
    basis, steps = [], []
    for k, v in enumerate(vecs):
        w = v.astype(float).copy()
        projs = []
        for j, b in enumerate(basis):
            coeff = float(np.dot(b, v) / np.dot(b, b))
            pv    = coeff * b
            projs.append({'j': j, 'coeff': coeff, 'vec': pv})
            w    -= pv
        nw  = float(np.linalg.norm(w))
        deg = nw < 1e-9
        u   = w / nw if not deg else np.zeros_like(w)
        steps.append({'k': k, 'v': v.copy(), 'w': w.copy(), 'u': u,
                      'projs': projs, 'norm_w': nw, 'degenerate': deg})
        if not deg:
            basis.append(u)
    return basis, steps


def gram_matrix_html(basis):
    """Compute the Gram matrix and output it as an HTML table."""
    n = len(basis)
    if n == 0:
        return '<p style="color:#888;font-size:12px;">No basis vectors yet</p>'

    rows_header = '<th></th>' + ''.join(
        f'<th style="padding:3px 10px;">e<sub>{j}</sub></th>' for j in range(n))

    rows = []
    all_ok = True
    for i in range(n):
        cells = [f'<td style="padding:3px 10px;"><b>e<sub>{i}</sub></b></td>']
        for j in range(n):
            val = float(np.dot(basis[i], basis[j]))
            if i == j:
                ok = abs(val - 1.0) < 1e-8
                expected = '1'
            else:
                ok = abs(val) < 1e-8
                expected = '0'
            if not ok:
                all_ok = False
            color  = '#0F6E56' if ok else '#D85A30'
            sym    = '✓' if ok else '✗'
            cells.append(
                f'<td style="padding:3px 10px;color:{color};font-weight:500;">'
                f'{val:.2e} {sym}</td>')
        rows.append('<tr>' + ''.join(cells) + '</tr>')

    status_color = '#0F6E56' if all_ok else '#888'
    status = 'All basis vectors are orthonormal ✓' if all_ok else 'Orthogonalization in progress…'

    return f"""
<div style='font-family:monospace;font-size:12px;background:#F8F8F8;
            border:1px solid #D6D6D6;border-radius:6px;padding:10px 14px;'>
  <div style='font-size:13px;font-weight:500;color:{status_color};margin-bottom:8px;'>
    Gram matrix &nbsp;G<sub>ij</sub> = ⟨e<sub>i</sub>, e<sub>j</sub>⟩ &nbsp;|&nbsp; {status}
  </div>
  <table style='border-collapse:collapse;'>
    <tr style='color:#888;font-size:11px;'>{rows_header}</tr>
    {''.join(rows)}
  </table>
</div>"""

# ============================================================
# Plotting in R²
# ============================================================

def arrow_annot(x0, y0, dx, dy, color, width=2.5, dash=False):
    """Plotly 2D vector annotation. dash=True → dashed arrow (input vector)."""
    # Plotly annotations do not support dashed arrows; the dashed effect is simulated with a Scatter segment + a small solid arrowhead
    # An annotation is still used here; the dash information is used only for the separate Scatter segment
    return dict(
        x=x0+dx, y=y0+dy, ax=x0, ay=y0,
        xref='x', yref='y', axref='x', ayref='y',
        showarrow=True,
        arrowhead=3, arrowsize=1.2, arrowwidth=width,
        arrowcolor=color,
    )


def plot_gs_2d(vecs, step_idx, show_basis):
    basis, steps = gram_schmidt_steps(vecs)
    n = len(vecs)
    traces, annots = [], []
    off = 0.12  # label offset

    # ── Input vectors (dashed segment + small solid arrowhead) ──
    for k, v in enumerate(vecs):
        col = VEC_INPUT_COLORS[k % len(VEC_INPUT_COLORS)]
        # dashed segment
        traces.append(go.Scatter(
            x=[0, v[0]*0.92], y=[0, v[1]*0.92],
            mode='lines',
            line=dict(color=col, width=2, dash='dash'),
            showlegend=False, hoverinfo='skip'
        ))
        # arrow tip (the annotation draws only the last short piece)
        annots.append(arrow_annot(v[0]*0.92, v[1]*0.92,
                                  v[0]*0.08, v[1]*0.08, col, width=2.0))
        # text label
        traces.append(go.Scatter(
            x=[v[0]+off], y=[v[1]+off],
            mode='text',
            text=[f'<b>v<sub>{k}</sub></b>'],
            textfont=dict(color=col, size=14),
            showlegend=False
        ))

    # ── Step-by-step projections and residuals ──
    for k in range(min(step_idx, n)):
        s = steps[k]
        if k == 0:
            continue  # vector 0 (v₀) is normalized directly; there is no projection step

        # projection vectors (from the origin, orange dotted line + arrow)
        for p in s['projs']:
            pv = p['vec']
            plen = np.linalg.norm(pv)
            if plen < 1e-10:
                continue
            traces.append(go.Scatter(
                x=[0, pv[0]*0.88], y=[0, pv[1]*0.88],
                mode='lines',
                line=dict(color=PROJ_COLOR, width=2, dash='dot'),
                showlegend=False, hoverinfo='skip'
            ))
            annots.append(arrow_annot(pv[0]*0.88, pv[1]*0.88,
                                      pv[0]*0.12, pv[1]*0.12,
                                      PROJ_COLOR, width=1.8))
            mid = pv * 0.5 + np.array([-0.1, 0.1])
            traces.append(go.Scatter(
                x=[mid[0]], y=[mid[1]],
                mode='text',
                text=[f'<span style="font-size:10px;">proj<sub>e{p["j"]}</sub></span>'],
                textfont=dict(color=PROJ_COLOR, size=10),
                showlegend=False
            ))

        # residual w (solid arrow)
        w   = s['w']
        nw  = s['norm_w']
        wcol = C_AUX if s['degenerate'] else RESIDUE_COLOR
        if nw > 1e-10:
            traces.append(go.Scatter(
                x=[0, w[0]*0.88], y=[0, w[1]*0.88],
                mode='lines', line=dict(color=wcol, width=2.5),
                showlegend=False, hoverinfo='skip'
            ))
            annots.append(arrow_annot(w[0]*0.88, w[1]*0.88,
                                      w[0]*0.12, w[1]*0.12, wcol, width=2.5))
        deg_note = ' ⚠ degenerate' if s['degenerate'] else f' ‖={nw:.3f}'
        traces.append(go.Scatter(
            x=[w[0]+off if nw>1e-10 else 0.1],
            y=[w[1]+off if nw>1e-10 else 0.1],
            mode='text',
            text=[f'<b>w<sub>{k}</sub></b>{deg_note}'],
            textfont=dict(color=wcol, size=11),
            showlegend=False
        ))

        # right-angle marks (w ⊥ the existing basis vectors)
        if not s['degenerate'] and nw > 1e-10:
            sq = 0.06
            for b in basis[:k]:
                bh = b / np.linalg.norm(b)
                wh = w / nw
                p1 = sq * bh;  p2 = sq*bh + sq*wh;  p3 = sq * wh
                traces.append(go.Scatter(
                    x=[p1[0], p2[0], p3[0]], y=[p1[1], p2[1], p3[1]],
                    mode='lines', line=dict(color=C_GRID, width=1),
                    showlegend=False
                ))

    # ── Final orthonormal basis (thick solid arrows) ──
    if show_basis:
        for j, b in enumerate(basis):
            col = VEC_BASIS_COLORS[j % len(VEC_BASIS_COLORS)]
            traces.append(go.Scatter(
                x=[0, b[0]*0.88], y=[0, b[1]*0.88],
                mode='lines', line=dict(color=col, width=3.5),
                showlegend=False, hoverinfo='skip'
            ))
            annots.append(arrow_annot(b[0]*0.88, b[1]*0.88,
                                      b[0]*0.12, b[1]*0.12, col, width=3.5))
            traces.append(go.Scatter(
                x=[b[0]+off], y=[b[1]+off],
                mode='text',
                text=[f'<b>e<sub>{j}</sub></b>'],
                textfont=dict(color=col, size=14),
                showlegend=False
            ))

    lim = max(max(np.linalg.norm(v) for v in vecs)*1.45, 2.0)
    fig = go.Figure(data=traces)
    fig.update_layout(
        annotations=annots,
        plot_bgcolor=C_BG, paper_bgcolor=C_BG,
        xaxis=dict(range=[-lim, lim], zeroline=True, zerolinecolor=C_AXIS,
                   gridcolor=C_GRID, gridwidth=0.8, scaleanchor='y',
                   showticklabels=True, title='x₀'),
        yaxis=dict(range=[-lim, lim], zeroline=True, zerolinecolor=C_AXIS,
                   gridcolor=C_GRID, gridwidth=0.8,
                   showticklabels=True, title='x₁'),
        showlegend=False, height=460,
        margin=dict(l=40, r=20, t=16, b=40),
    )
    return fig, basis, steps

# ============================================================
# Plotting in R³ (focusing on orthogonality, without listing the projection details)
# ============================================================

def add_vec_3d(fig, end, color, name, dashed=False, width=5):
    """Add a vector to the 3D scene (shaft + cone head)."""
    end = np.asarray(end, float)
    length = np.linalg.norm(end)
    if length < 1e-12:
        return
    ratio     = 0.15
    shaft_end = end * (1 - ratio)
    dash_kw   = dict(dash='dash') if dashed else {}

    fig.add_trace(go.Scatter3d(
        x=[0, shaft_end[0]], y=[0, shaft_end[1]], z=[0, shaft_end[2]],
        mode='lines',
        line=dict(color=color, width=width, **dash_kw),
        name=name, showlegend=True, legendgroup=name
    ))
    cone_sz = length * ratio * 0.7
    fig.add_trace(go.Cone(
        x=[end[0]], y=[end[1]], z=[end[2]],
        u=[end[0]-shaft_end[0]], v=[end[1]-shaft_end[1]], w=[end[2]-shaft_end[2]],
        sizemode='absolute', sizeref=cone_sz,
        colorscale=[[0, color],[1, color]],
        showscale=False,
        name=name, showlegend=False, legendgroup=name
    ))


def plot_gs_3d(vecs, step_idx, show_basis):
    """
    Visualization strategy in R³:
    - always show the input vectors (dashed)
    - show the computed orthonormal basis vectors step by step (thick solid arrows)
    - do not list every projection vector (too cluttered in 3D)
    - the Gram matrix is updated in real time in the side panel
    """
    basis, steps = gram_schmidt_steps(vecs)
    n = len(vecs)
    fig = go.Figure()

    # ── Input vectors (dashed) ──
    for k, v in enumerate(vecs):
        col = VEC_INPUT_COLORS[k % len(VEC_INPUT_COLORS)]
        add_vec_3d(fig, v, col, f'v{k} (input)', dashed=True, width=4)
        # label
        fig.add_trace(go.Scatter3d(
            x=[v[0]*1.1], y=[v[1]*1.1], z=[v[2]*1.1],
            mode='text', text=[f'v{k}'],
            textfont=dict(color=col, size=13),
            showlegend=False
        ))

    # ── Show the residuals w_k step by step (not controlled by show_basis) ──
    for k in range(1, min(step_idx, n)):
        s = steps[k]
        w   = s['w']
        nw  = s['norm_w']
        if s['degenerate']:
            wcol = C_AUX
            label = f'w{k} ⚠ degenerate'
        else:
            wcol  = RESIDUE_COLOR
            label = f'w{k}  ‖={nw:.3f}'
        if nw > 1e-10:
            add_vec_3d(fig, w, wcol, label, dashed=False, width=5)

    # ── Show the orthonormal basis vectors (only when show_basis is checked, up to step_idx) ──
    if show_basis:
        for k in range(min(step_idx, n)):
            s = steps[k]
            if s['degenerate']:
                continue
            col = VEC_BASIS_COLORS[k % len(VEC_BASIS_COLORS)]
            add_vec_3d(fig, s['u'], col, f'e{k}', dashed=False, width=7)
            b = s['u']
            fig.add_trace(go.Scatter3d(
                x=[b[0]*1.1], y=[b[1]*1.1], z=[b[2]*1.1],
                mode='text', text=[f'e{k}'],
                textfont=dict(color=col, size=14),
                showlegend=False
            ))

    lim = max(max(np.linalg.norm(v) for v in vecs)*1.35, 1.5)
    fig.update_layout(
        scene=dict(
            bgcolor=C_BG,
            xaxis=dict(range=[-lim,lim], gridcolor=C_GRID,
                       backgroundcolor=C_BG, title='x₀'),
            yaxis=dict(range=[-lim,lim], gridcolor=C_GRID,
                       backgroundcolor=C_BG, title='x₁'),
            zaxis=dict(range=[-lim,lim], gridcolor=C_GRID,
                       backgroundcolor=C_BG, title='x₂'),
            aspectmode='cube',
        ),
        showlegend=True,
        legend=dict(x=0.0, y=1.0, bgcolor='rgba(0,0,0,0)', font=dict(size=12)),
        height=480,
        margin=dict(l=0, r=0, t=16, b=0),
        paper_bgcolor=C_BG,
    )
    return fig, basis, steps

# ============================================================
# Interface widgets and update logic
# ============================================================

# ---------- Mode ----------
mode_tb = widgets.ToggleButtons(
    options=[('ℝ² plane', '2d'), ('ℝ³ space', '3d')],
    description='Dimension:',
    style={'description_width':'60px'},
    layout=widgets.Layout(width='300px')
)

step_sl = widgets.IntSlider(
    value=0, min=0, max=2, step=1,
    description='Processed to v:',
    continuous_update=True,
    style={'description_width':'90px'},
    layout=widgets.Layout(width='360px')
)

show_chk = widgets.Checkbox(
    value=False, description='Show orthonormal basis e_k',
    layout=widgets.Layout(width='200px')
)

# ---------- Vector input (R²) ----------
def ft(v, desc):
    return widgets.BoundedFloatText(
        value=v, min=-5.0, max=5.0, step=0.1,
        description=desc,
        style={'description_width':'24px'},
        layout=widgets.Layout(width='82px')
    )

r2_00=ft(1.0,'x₀'); r2_01=ft(1.0,'y₀')
r2_10=ft(1.0,'x₁'); r2_11=ft(0.0,'y₁')

r2_row = widgets.HBox([
    widgets.HTML('<b style="font-size:12px;line-height:32px;">v₀&nbsp;</b>'),
    r2_00, r2_01,
    widgets.HTML('<b style="font-size:12px;line-height:32px;">&nbsp;&nbsp;v₁&nbsp;</b>'),
    r2_10, r2_11,
])

# ---------- Vector input (R³) ----------
r3_0 = [ft(1.0,'x₀'), ft(1.0,'y₀'), ft(0.0,'z₀')]
r3_1 = [ft(1.0,'x₁'), ft(0.0,'y₁'), ft(1.0,'z₁')]
r3_2 = [ft(0.0,'x₂'), ft(1.0,'y₂'), ft(1.0,'z₂')]

def vec_row(label, widgets_list):
    return widgets.HBox([
        widgets.HTML(f'<b style="font-size:12px;line-height:32px;">{label}&nbsp;</b>')
    ] + widgets_list)

r3_rows = widgets.VBox([
    vec_row('v₀', r3_0),
    vec_row('v₁', r3_1),
    vec_row('v₂', r3_2),
])

# ---------- Preset buttons ----------
def mkbtn(label, color='#F0EAF7'):
    return widgets.Button(
        description=label,
        style={'button_color': color},
        layout=widgets.Layout(width='auto', height='28px')
    )

btn_reset   = mkbtn('Reset')
btn_nondep  = mkbtn('⚠ Nearly dependent', '#FFF0EC')
btn_orth    = mkbtn('Orthogonal set', '#EAF9F8')

# ---------- Outputs ----------
out_fig  = widgets.Output()
out_step = widgets.HTML(value='')
out_gram = widgets.HTML(value='')


def get_vecs():
    if mode_tb.value == '2d':
        return [np.array([r2_00.value, r2_01.value]),
                np.array([r2_10.value, r2_11.value])]
    else:
        return [np.array([w.value for w in r3_0]),
                np.array([w.value for w in r3_1]),
                np.array([w.value for w in r3_2])]


def step_desc(step_idx, steps):
    """Generate the text description of the current step."""
    n = len(steps)
    if step_idx == 0:
        return ('<span style="color:#555;font-size:13px;">'
                'Step 0: the input vectors are shown; orthogonalization has not started yet.</span>')
    lines = []
    for k in range(min(step_idx, n)):
        s = steps[k]
        if k == 0:
            nv = np.linalg.norm(s['v'])
            lines.append(
                f'<b>Step 1</b>: u₀ = v₀, normalize → '
                f'<b style="color:{VEC_BASIS_COLORS[0]};">e₀</b> = v₀ / {nv:.3f}'
            )
        else:
            psum = ' + '.join(
                f'proj<sub>e{p["j"]}</sub>(v<sub>{k}</sub>)'
                for p in s['projs']
            ) or '0'
            if s['degenerate']:
                note = (f'<span style="color:#D85A30;">'
                        f'⚠ ‖w<sub>{k}</sub>‖ ≈ 0 — v<sub>{k}</sub> is nearly linearly dependent; skipped</span>')
            else:
                col = VEC_BASIS_COLORS[k % len(VEC_BASIS_COLORS)]
                note = (f'‖w<sub>{k}</sub>‖ = {s["norm_w"]:.4f}, '
                        f'normalize → <b style="color:{col};">e<sub>{k}</sub></b>')
            lines.append(
                f'<b>Step {k+1}</b>: '
                f'w<sub>{k}</sub> = v<sub>{k}</sub> − ({psum}) &nbsp;→&nbsp; {note}'
            )
    return '<br>'.join(lines)


def update(_=None):
    vecs = get_vecs()
    mode = mode_tb.value
    si   = step_sl.value
    sb   = show_chk.value

    # keep the slider maximum in sync
    step_sl.max = len(vecs)

    if mode == '2d':
        fig, basis, steps = plot_gs_2d(vecs, si, sb)
    else:
        fig, basis, steps = plot_gs_3d(vecs, si, sb)

    with out_fig:
        out_fig.clear_output(wait=True)
        fig.show()

    out_step.value = (
        '<div style="background:#F8F8F8;border:1px solid #D6D6D6;'
        'border-radius:6px;padding:8px 12px;font-size:13px;line-height:1.8;">'
        + step_desc(si, steps) + '</div>'
    )
    # Gram matrix: show the basis only up to the current step
    _, partial_steps = gram_schmidt_steps(vecs)
    partial_basis = []
    for k in range(min(si, len(partial_steps))):
        s = partial_steps[k]
        if not s['degenerate']:
            partial_basis.append(s['u'])
    out_gram.value = gram_matrix_html(partial_basis)


# ---------- Preset button logic ----------
def set_r2(v0, v1):
    r2_00.value, r2_01.value = v0
    r2_10.value, r2_11.value = v1

def set_r3(v0, v1, v2):
    for w, val in zip(r3_0, v0): w.value = val
    for w, val in zip(r3_1, v1): w.value = val
    for w, val in zip(r3_2, v2): w.value = val

def on_reset(_):
    if mode_tb.value == '2d': set_r2((1.,1.), (1.,0.))
    else: set_r3((1.,1.,0.), (1.,0.,1.), (0.,1.,1.))
    update()

def on_nondep(_):
    if mode_tb.value == '2d': set_r2((1.,1.), (1.05, 1.0))
    else: set_r3((1.,0.,0.), (1.0, 0.08, 0.), (0.,1.,1.))
    update()

def on_orth(_):
    if mode_tb.value == '2d': set_r2((1.,0.), (0.,2.))
    else: set_r3((2.,0.,0.), (0.,3.,0.), (0.,0.,1.5))
    update()

btn_reset.on_click(on_reset)
btn_nondep.on_click(on_nondep)
btn_orth.on_click(on_orth)

# ---------- Observers ----------
for w in [mode_tb, step_sl, show_chk,
          r2_00, r2_01, r2_10, r2_11,
          *r3_0, *r3_1, *r3_2]:
    w.observe(update, names='value')

def on_mode(_=None):
    is2 = mode_tb.value == '2d'
    r2_row.layout.display   = '' if is2 else 'none'
    r3_rows.layout.display  = 'none' if is2 else ''
    step_sl.max   = 2 if is2 else 3
    step_sl.value = 0
    update()

mode_tb.observe(on_mode, names='value')
r3_rows.layout.display = 'none'

# ============================================================
# Assembling the layout
# ============================================================

header = widgets.HTML(
    '<h3 style="margin:4px 0 2px;color:#57068C;font-family:sans-serif;">'
    'Gram-Schmidt Orthogonalization — Step-by-Step Visualization</h3>'
    '<p style="margin:0 0 10px;font-size:12px;color:#666;">'
    'dashed = input vectors v_k &nbsp;|&nbsp; solid = orthonormal basis e_k &nbsp;|&nbsp; '
    'orange arrows = projections proj (ℝ² only)</p>'
)

ctrl_row = widgets.HBox(
    [mode_tb, widgets.VBox([step_sl, show_chk])],
    layout=widgets.Layout(gap='20px', align_items='flex-start')
)

preset_row = widgets.HBox(
    [widgets.HTML('<span style="font-size:12px;line-height:28px;">Presets:</span>'),
     btn_reset, btn_nondep, btn_orth],
    layout=widgets.Layout(gap='8px', align_items='center', margin='4px 0')
)

gram_box = widgets.VBox(
    [widgets.HTML('<b style="font-size:13px;">Orthogonality check</b>'), out_gram],
    layout=widgets.Layout(margin='8px 0')
)

ui = widgets.VBox([
    header, ctrl_row,
    widgets.HTML('<hr style="margin:8px 0;border-color:#D6D6D6;">'),
    widgets.HTML('<b style="font-size:12px;">Input vectors (range −5 to 5)</b>'),
    r2_row, r3_rows, preset_row,
    widgets.HTML('<hr style="margin:8px 0;border-color:#D6D6D6;">'),
    out_fig,
    out_step,
    gram_box,
])

display(ui)
update()