feat(cholesky): improve layout and controls for Cholesky decomposition and solving linear systems

pages/linear_algebra/02_cholesky.py

@@ -4,20 +4,20 @@ import numpy as np
 
 st.title("Cholesky Decomposition")
 
-st.markdown(
-    r"""
-For a **symmetric positive definite** matrix $A$, Cholesky gives $A = LL^\top$ 
-where $L$ is lower triangular. The algorithm computes $L$ element by element.
-"""
-)
-
-# Size selection
-col_size, col_preset = st.columns(2)
-
-with col_size:
+# st.markdown(
+#     r"""
+# For a **symmetric positive definite** matrix $A$, Cholesky gives $A = LL^\top$
+# where $L$ is lower triangular. The algorithm computes $L$ element by element.
+# """
+# )
+
+# Controls and visualizations layout
+col_controls, col_content = st.columns([1, 3])
+
+# LEFT COLUMN: Controls
+with col_controls:
     n = st.radio("Matrix size", [3, 4], horizontal=True)
 
-with col_preset:
     preset = st.selectbox(
         "Example matrix",
         ["Simple diagonal-dominant", "Correlation matrix", "Random SPD"],
@@ -111,30 +111,6 @@ def cholesky_steps(A):
 steps = cholesky_steps(A)
 total_steps = len(steps)
 
-# Buttons
-col_prev, col_next, col_reset = st.columns([1, 1, 1])
-
-with col_prev:
-    if st.button("← Previous", disabled=st.session_state.chol_step == 0):
-        st.session_state.chol_step -= 1
-        st.rerun()
-
-with col_next:
-    if st.button("Next →", disabled=st.session_state.chol_step >= total_steps):
-        st.session_state.chol_step += 1
-        st.rerun()
-
-with col_reset:
-    if st.button("Reset"):
-        st.session_state.chol_step = 0
-        st.rerun()
-
-current_step = st.session_state.chol_step
-
-# Progress
-st.progress(current_step / total_steps if total_steps > 0 else 0)
-st.caption(f"Step {current_step} of {total_steps}")
-
 
 # Helper to format matrix as LaTeX
 def matrix_to_latex(M, highlight=None, computed_mask=None):
@@ -160,93 +136,129 @@ def matrix_to_latex(M, highlight=None, computed_mask=None):
     return r"\begin{pmatrix} " + r" \\ ".join(rows) + r" \end{pmatrix}"
 
 
-# Display matrices side by side
-st.markdown("---")
-
-col_A, col_eq, col_L, col_Lt = st.columns([2, 0.5, 2, 2])
-
-# Get current L state
-if current_step == 0:
-    L_current = np.zeros_like(A)
-    computed = np.zeros((n, n), dtype=bool)
-    highlight = None
-else:
-    i, j, formula, val, L_current = steps[current_step - 1]
-    computed = np.abs(L_current) > 1e-10
-    # Upper triangle is always "computed" as 0
-    for ii in range(n):
-        for jj in range(ii + 1, n):
-            computed[ii, jj] = True
-    highlight = (i, j)
-
-with col_A:
-    st.markdown("**Matrix $A$:**")
-    st.latex(f"A = {matrix_to_latex(A)}")
-
-with col_eq:
-    st.markdown(" ")
-    st.latex("=")
-
-with col_L:
-    st.markdown("**Cholesky $L$:**")
-    # Create mask for computed entries
-    mask = np.zeros((n, n), dtype=bool)
-    for ii in range(n):
-        for jj in range(ii + 1, n):
-            mask[ii, jj] = True  # Upper triangle always shown as 0
-    if current_step > 0:
-        for step_idx in range(current_step):
-            si, sj, _, _, _ = steps[step_idx]
-            mask[si, sj] = True
-    st.latex(
-        f"L = {matrix_to_latex(L_current, highlight=highlight, computed_mask=mask)}"
-    )
+# Add buttons and progress to left column
+with col_controls:
+    st.markdown("---")
 
-with col_Lt:
-    st.markdown("**Transpose $L^\\top$:**")
-    L_T = L_current.T
-    mask_T = mask.T
-    highlight_T = (highlight[1], highlight[0]) if highlight else None
-    st.latex(
-        f"L^\\top = {matrix_to_latex(L_T, highlight=highlight_T, computed_mask=mask_T)}"
-    )
+    # Buttons
+    col_prev, col_next, col_reset = st.columns([1, 1, 1])
+
+    with col_prev:
+        if st.button(
+            "← Prev", disabled=st.session_state.chol_step == 0, use_container_width=True
+        ):
+            st.session_state.chol_step -= 1
+            st.rerun()
+
+    with col_next:
+        if st.button(
+            "Next →",
+            disabled=st.session_state.chol_step >= total_steps,
+            use_container_width=True,
+        ):
+            st.session_state.chol_step += 1
+            st.rerun()
+
+    with col_reset:
+        if st.button("Reset", use_container_width=True):
+            st.session_state.chol_step = 0
+            st.rerun()
+
+    # Progress
+    st.progress(st.session_state.chol_step / total_steps if total_steps > 0 else 0)
+    st.caption(f"Step {st.session_state.chol_step} of {total_steps}")
+
+# RIGHT COLUMN: Content
+with col_content:
+    current_step = st.session_state.chol_step
+
+    # Display matrices side by side
+    st.markdown("---")
 
-# Show current step formula
-st.markdown("---")
+    col_A, col_eq, col_L, col_Lt = st.columns([2, 0.5, 2, 2])
 
-if current_step == 0:
-    st.markdown(
-        "**Cholesky Algorithm:** For each column $j$, compute elements from top to bottom:"
-    )
-    st.latex(
-        r"L_{jj} = \sqrt{A_{jj} - \sum_{k=1}^{j-1} L_{jk}^2} \qquad L_{ij} = \frac{A_{ij} - \sum_{k=1}^{j-1} L_{ik} L_{jk}}{L_{jj}} \text{ for } i > j"
-    )
-    st.info("Click **Next →** to start computing elements of $L$")
-elif current_step <= total_steps:
-    i, j, formula, val, _ = steps[current_step - 1]
-    if i == j:
-        st.markdown(f"**Computing diagonal element $L_{{{i + 1}{j + 1}}}$:**")
+    # Get current L state
+    if current_step == 0:
+        L_current = np.zeros_like(A)
+        computed = np.zeros((n, n), dtype=bool)
+        highlight = None
     else:
-        st.markdown(f"**Computing off-diagonal element $L_{{{i + 1}{j + 1}}}$:**")
-    st.latex(formula + f" = {val:.4f}")
+        i, j, formula, val, L_current = steps[current_step - 1]
+        computed = np.abs(L_current) > 1e-10
+        # Upper triangle is always "computed" as 0
+        for ii in range(n):
+            for jj in range(ii + 1, n):
+                computed[ii, jj] = True
+        highlight = (i, j)
+
+    with col_A:
+        st.markdown("**Matrix $A$:**")
+        st.latex(f"A = {matrix_to_latex(A)}")
 
-# Verification at the end
-if current_step == total_steps:
-    st.markdown("---")
-    st.success("**✓ Cholesky decomposition complete!**")
-    L_final = steps[-1][4]
-    LLT = L_final @ L_final.T
-    st.markdown("**Verification:** $LL^\\top = A$")
-
-    col_v1, col_v2, col_v3 = st.columns([2, 0.5, 2])
-    with col_v1:
-        st.latex(f"LL^\\top = {matrix_to_latex(LLT)}")
-    with col_v2:
+    with col_eq:
+        st.markdown("&nbsp;")
         st.latex("=")
-    with col_v3:
-        st.latex(f"A = {matrix_to_latex(A)}")
 
-    # Check numerical accuracy
-    error = np.max(np.abs(LLT - A))
-    if error < 1e-10:
-        st.markdown(f"Maximum error: ${error:.2e}$ ✓")
+    with col_L:
+        st.markdown("**Cholesky $L$:**")
+        # Create mask for computed entries
+        mask = np.zeros((n, n), dtype=bool)
+        for ii in range(n):
+            for jj in range(ii + 1, n):
+                mask[ii, jj] = True  # Upper triangle always shown as 0
+        if current_step > 0:
+            for step_idx in range(current_step):
+                si, sj, _, _, _ = steps[step_idx]
+                mask[si, sj] = True
+        st.latex(
+            f"L = {matrix_to_latex(L_current, highlight=highlight, computed_mask=mask)}"
+        )
+
+    with col_Lt:
+        st.markdown("**Transpose $L^\\top$:**")
+        L_T = L_current.T
+        mask_T = mask.T
+        highlight_T = (highlight[1], highlight[0]) if highlight else None
+        st.latex(
+            f"L^\\top = {matrix_to_latex(L_T, highlight=highlight_T, computed_mask=mask_T)}"
+        )
+
+    # Show current step formula
+    st.markdown("---")
+
+    if current_step == 0:
+        st.markdown(
+            "**Cholesky Algorithm:** For each column $j$, compute elements from top to bottom:"
+        )
+        st.latex(
+            r"L_{jj} = \sqrt{A_{jj} - \sum_{k=1}^{j-1} L_{jk}^2} \qquad L_{ij} = \frac{A_{ij} - \sum_{k=1}^{j-1} L_{ik} L_{jk}}{L_{jj}} \text{ for } i > j"
+        )
+        st.info("Click **Next →** to start computing elements of $L$")
+    elif current_step <= total_steps:
+        i, j, formula, val, _ = steps[current_step - 1]
+        if i == j:
+            st.markdown(f"**Computing diagonal element $L_{{{i + 1}{j + 1}}}$:**")
+        else:
+            st.markdown(f"**Computing off-diagonal element $L_{{{i + 1}{j + 1}}}$:**")
+        st.latex(formula + f" = {val:.4f}")
+
+    # Verification at the end
+    if current_step == total_steps:
+        st.markdown("---")
+        st.success("**✓ Cholesky decomposition complete!**")
+        # L_final = steps[-1][4]
+        # LLT = L_final @ L_final.T
+        # st.markdown("**Verification:** $LL^\\top = A$")
+
+        # col_v1, col_v2, col_v3 = st.columns([2, 0.5, 2])
+        # with col_v1:
+        #     st.latex(f"LL^\\top = {matrix_to_latex(LLT)}")
+        # with col_v2:
+        #     st.latex("=")
+        # with col_v3:
+        #     st.latex(f"A = {matrix_to_latex(A)}")
+
+        # # Check numerical accuracy
+        # error = np.max(np.abs(LLT - A))
+        # if error < 1e-10:
+        #     st.markdown(f"Maximum error: ${error:.2e}$ ✓")

pages/linear_algebra/03_cholesky_solve.py

@@ -4,22 +4,12 @@ import numpy as np
 
 st.title("Solving Linear Systems with Cholesky")
 
-st.markdown(
-    r"""
-To solve $A\mathbf{x} = \mathbf{b}$ where $A$ is symmetric positive definite:
-1. Compute Cholesky: $A = LL^\top$
-2. Solve $L\mathbf{y} = \mathbf{b}$ (forward substitution)
-3. Solve $L^\top\mathbf{x} = \mathbf{y}$ (backward substitution)
-"""
-)
-
-# Size selection
-col_size, col_preset = st.columns(2)
-
-with col_size:
-    n = st.radio("System size", [3, 4], horizontal=True)
+# Create layout
+col_controls, col_content = st.columns([1, 3])
 
-with col_preset:
+# LEFT COLUMN: Controls
+with col_controls:
+    n = st.radio("System size", [3, 4], horizontal=True)
     preset = st.selectbox(
         "Example",
         ["Simple system", "Unit RHS", "Random"],
@@ -29,8 +19,8 @@ with col_preset:
 np.random.seed(123)
 if preset == "Simple system":
     if n == 3:
-        A = np.array([[4.0, 2.0, 1.0], [2.0, 5.0, 2.0], [1.0, 2.0, 6.0]])
-        b = np.array([1.0, 2.0, 3.0])
+        A = np.array([[1.0, 1.0, 1.0], [1.0, 2.0, 2.0], [1.0, 2.0, 3.0]])
+        b = np.array([3.0, 5.0, 6.0])
     else:
         A = np.array(
             [
@@ -138,46 +128,58 @@ if st.session_state.solve_n != n or st.session_state.solve_preset != preset:
     st.session_state.solve_n = n
     st.session_state.solve_preset = preset
 
-# Buttons
-col_prev, col_next, col_reset = st.columns([1, 1, 1])
-
-with col_prev:
-    if st.button("← Previous", disabled=st.session_state.solve_step == 0):
-        st.session_state.solve_step -= 1
-        st.rerun()
-
-with col_next:
-    if st.button("Next →", disabled=st.session_state.solve_step >= total_steps):
-        st.session_state.solve_step += 1
-        st.rerun()
+# Continue LEFT COLUMN: Add buttons and progress
+with col_controls:
+    st.markdown("---")
 
-with col_reset:
-    if st.button("Reset"):
-        st.session_state.solve_step = 0
-        st.rerun()
+    # Buttons
+    col_prev, col_next, col_reset = st.columns([1, 1, 1])
+
+    with col_prev:
+        if st.button(
+            "← Prev",
+            disabled=st.session_state.solve_step == 0,
+            use_container_width=True,
+        ):
+            st.session_state.solve_step -= 1
+            st.rerun()
+
+    with col_next:
+        if st.button(
+            "Next →",
+            disabled=st.session_state.solve_step >= total_steps,
+            use_container_width=True,
+        ):
+            st.session_state.solve_step += 1
+            st.rerun()
+
+    with col_reset:
+        if st.button("Reset", use_container_width=True):
+            st.session_state.solve_step = 0
+            st.rerun()
+
+    # Progress
+    st.progress(st.session_state.solve_step / total_steps if total_steps > 0 else 0)
+
+    # Caption
+    phase = (
+        "setup"
+        if st.session_state.solve_step == 0
+        else ("forward" if st.session_state.solve_step <= total_forward else "backward")
+    )
+    if phase == "setup":
+        st.caption("Step 0: Setup")
+    elif phase == "forward":
+        st.caption(
+            f"Forward substitution: Step {st.session_state.solve_step} of {total_forward}"
+        )
+    else:
+        st.caption(
+            f"Backward substitution: Step {st.session_state.solve_step - total_forward} of {total_backward}"
+        )
 
 current_step = st.session_state.solve_step
 
-# Progress
-st.progress(current_step / total_steps if total_steps > 0 else 0)
-
-# Determine phase
-if current_step == 0:
-    phase = "setup"
-elif current_step <= total_forward:
-    phase = "forward"
-else:
-    phase = "backward"
-
-if phase == "setup":
-    st.caption("Step 0: Setup")
-elif phase == "forward":
-    st.caption(f"Forward substitution: Step {current_step} of {total_forward}")
-else:
-    st.caption(
-        f"Backward substitution: Step {current_step - total_forward} of {total_backward}"
-    )
-
 
 # Helper to format vector as LaTeX
 def vector_to_latex(v, highlight=None, computed_mask=None):
@@ -204,115 +206,96 @@ def matrix_to_latex(M):
     return r"\begin{pmatrix} " + r" \\ ".join(rows) + r" \end{pmatrix}"
 
 
-# Display system
-st.markdown("---")
-
-# Show the problem setup
-col1, col2, col3 = st.columns([3, 3, 3])
-
-with col1:
-    st.markdown("**System $A\\mathbf{x} = \\mathbf{b}$:**")
-    st.latex(f"A = {matrix_to_latex(A)}")
-    st.latex(f"\\mathbf{{b}} = {vector_to_latex(b)}")
-
-with col2:
-    st.markdown("**Cholesky factor $L$:**")
-    st.latex(f"L = {matrix_to_latex(L)}")
-
-# Compute current state of y and x
-if current_step == 0:
-    y_current = np.zeros(n)
-    x_current = np.zeros(n)
-    y_mask = np.zeros(n, dtype=bool)
-    x_mask = np.zeros(n, dtype=bool)
-    y_highlight = None
-    x_highlight = None
-elif current_step <= total_forward:
-    # In forward substitution
-    _, _, _, y_current = forward_steps[current_step - 1]
-    x_current = np.zeros(n)
-    y_mask = np.zeros(n, dtype=bool)
-    for s in range(current_step):
-        idx = forward_steps[s][0]
-        y_mask[idx] = True
-    x_mask = np.zeros(n, dtype=bool)
-    y_highlight = forward_steps[current_step - 1][0]
-    x_highlight = None
-else:
-    # In backward substitution
-    y_current = y_final.copy()
-    backward_idx = current_step - total_forward - 1
-    _, _, _, x_current = backward_steps[backward_idx]
-    y_mask = np.ones(n, dtype=bool)
-    x_mask = np.zeros(n, dtype=bool)
-    for s in range(backward_idx + 1):
-        idx = backward_steps[s][0]
-        x_mask[idx] = True
-    y_highlight = None
-    x_highlight = backward_steps[backward_idx][0]
-
-with col3:
-    st.markdown("**Intermediate $\\mathbf{y}$ and solution $\\mathbf{x}$:**")
-    col_y, col_x = st.columns(2)
-    with col_y:
-        st.latex(
-            f"\\mathbf{{y}} = {vector_to_latex(y_current, highlight=y_highlight, computed_mask=y_mask)}"
-        )
-    with col_x:
-        st.latex(
-            f"\\mathbf{{x}} = {vector_to_latex(x_current, highlight=x_highlight, computed_mask=x_mask)}"
-        )
+# RIGHT COLUMN: Content
+with col_content:
+    # Show the problem setup
+    col1, col2, col3 = st.columns([3, 3, 3])
 
-# Show current computation
-st.markdown("---")
+    with col1:
+        st.markdown("**System $A\\mathbf{x} = \\mathbf{b}$:**")
+        st.latex(f"A = {matrix_to_latex(A)}")
+        st.latex(f"\\mathbf{{b}} = {vector_to_latex(b)}")
 
-if phase == "setup":
-    st.markdown("**Step 1: Forward Substitution** — Solve $L\\mathbf{y} = \\mathbf{b}$")
-    st.latex(f"{matrix_to_latex(L)} \\mathbf{{y}} = {vector_to_latex(b)}")
-    st.markdown(
-        "Start from the top row and work down. Each $y_i$ depends only on previously computed values."
-    )
-    st.info("Click **Next →** to begin forward substitution")
+    with col2:
+        st.markdown("**Cholesky factor $L$:**")
+        st.latex(f"L = {matrix_to_latex(L)}")
+
+    # Compute current state of y and x
+    if current_step == 0:
+        y_current = np.zeros(n)
+        x_current = np.zeros(n)
+        y_mask = np.zeros(n, dtype=bool)
+        x_mask = np.zeros(n, dtype=bool)
+        y_highlight = None
+        x_highlight = None
+    elif current_step <= total_forward:
+        # In forward substitution
+        _, _, _, y_current = forward_steps[current_step - 1]
+        x_current = np.zeros(n)
+        y_mask = np.zeros(n, dtype=bool)
+        for s in range(current_step):
+            idx = forward_steps[s][0]
+            y_mask[idx] = True
+        x_mask = np.zeros(n, dtype=bool)
+        y_highlight = forward_steps[current_step - 1][0]
+        x_highlight = None
+    else:
+        # In backward substitution
+        y_current = y_final.copy()
+        backward_idx = current_step - total_forward - 1
+        _, _, _, x_current = backward_steps[backward_idx]
+        y_mask = np.ones(n, dtype=bool)
+        x_mask = np.zeros(n, dtype=bool)
+        for s in range(backward_idx + 1):
+            idx = backward_steps[s][0]
+            x_mask[idx] = True
+        y_highlight = None
+        x_highlight = backward_steps[backward_idx][0]
+
+    with col3:
+        st.markdown("**Intermediate $\\mathbf{y}$ and solution $\\mathbf{x}$:**")
+        col_y, col_x = st.columns(2)
+        with col_y:
+            st.latex(
+                f"\\mathbf{{y}} = {vector_to_latex(y_current, highlight=y_highlight, computed_mask=y_mask)}"
+            )
+        with col_x:
+            st.latex(
+                f"\\mathbf{{x}} = {vector_to_latex(x_current, highlight=x_highlight, computed_mask=x_mask)}"
+            )
 
-elif phase == "forward":
-    step_idx = current_step - 1
-    i, formula, val, _ = forward_steps[step_idx]
-    st.markdown(f"**Forward substitution:** Computing $y_{{{i + 1}}}$")
-    st.latex(formula + f" = {val:.4f}")
+    # Show current computation
+    st.markdown("---")
 
-    if current_step == total_forward:
-        st.success("✓ Forward substitution complete!")
+    if phase == "setup":
         st.markdown(
-            "**Next:** Backward substitution to solve $L^\\top \\mathbf{x} = \\mathbf{y}$"
+            "**Step 1: Forward Substitution** — Solve $L\\mathbf{y} = \\mathbf{b}$"
         )
+        st.latex(f"{matrix_to_latex(L)} \\mathbf{{y}} = {vector_to_latex(b)}")
+        st.markdown(
+            "Start from the top row and work down. Each $y_i$ depends only on previously computed values."
+        )
+        st.info("Click **Next →** to begin forward substitution")
+
+    elif phase == "forward":
+        step_idx = current_step - 1
+        i, formula, val, _ = forward_steps[step_idx]
+        st.markdown(f"**Forward substitution:** Computing $y_{{{i + 1}}}$")
+        st.latex(formula + f" = {val:.4f}")
+
+        if current_step == total_forward:
+            st.success("✓ Forward substitution complete!")
+            st.markdown(
+                "**Next:** Backward substitution to solve $L^\\top \\mathbf{x} = \\mathbf{y}$"
+            )
 
-else:  # backward
-    step_idx = current_step - total_forward - 1
-    i, formula, val, _ = backward_steps[step_idx]
-    st.markdown(f"**Backward substitution:** Computing $x_{{{i + 1}}}$")
-    st.latex(formula + f" = {val:.4f}")
-
-# Final verification
-if current_step == total_steps:
-    st.markdown("---")
-    st.success("**✓ System solved!**")
-
-    x_final = backward_steps[-1][3]
-    Ax = A @ x_final
-
-    st.markdown("**Verification:** $A\\mathbf{x} = \\mathbf{b}$")
-
-    col_v1, col_v2, col_v3, col_v4, col_v5 = st.columns([2, 0.3, 1.5, 0.3, 1.5])
-    with col_v1:
-        st.latex(f"A\\mathbf{{x}} = {matrix_to_latex(A)} {vector_to_latex(x_final)}")
-    with col_v2:
-        st.latex("=")
-    with col_v3:
-        st.latex(f"{vector_to_latex(Ax)}")
-    with col_v4:
-        st.latex("=")
-    with col_v5:
-        st.latex(f"\\mathbf{{b}} = {vector_to_latex(b)}")
+    else:  # backward
+        step_idx = current_step - total_forward - 1
+        i, formula, val, _ = backward_steps[step_idx]
+        st.markdown(f"**Backward substitution:** Computing $x_{{{i + 1}}}$")
+        st.latex(formula + f" = {val:.4f}")
 
-    error = np.max(np.abs(Ax - b))
-    st.markdown(f"Maximum error: ${error:.2e}$ ✓")
+    # Final verification
+    if current_step == total_steps:
+        st.markdown("---")
+        st.success("**✓ System solved!**")