""" Operator Splitting Explainer — Animation for @fminxyz channel ============================================================= Visualizes Lie-Trotter and Strang splitting vs exact solution on a 2D ODE: du/dt = A*u + B*u where: A = rotation (antisymmetric matrix) — the "interesting hard part" B = exponential decay (diagonal) — the "other part" Exact solution: matrix exponential of (A+B) Lie-Trotter: exp(A*dt) * exp(B*dt) applied sequentially — 1st order Strang: exp(A*dt/2) * exp(B*dt) * exp(A*dt/2) — 2nd order Output: splitting_animation.mp4, 1080x1080, ~30 frames, fps=2 """ import numpy as np from scipy.linalg import expm from scipy.integrate import solve_ivp import matplotlib.pyplot as plt import matplotlib.animation as animation from matplotlib.patches import FancyArrowPatch from matplotlib.lines import Line2D # ── Palette & style ───────────────────────────────────────────────────────── BG = "#1a1a2e" PANEL = "#16213e" ACCENT = "#0f3460" EXACT = "#e94560" # red — exact LT_COL = "#f5a623" # orange — Lie-Trotter STRANG = "#7bc8f6" # sky blue — Strang GRID_COL = "#2a2a4a" WHITE = "#e8e8f0" DIM = "#6a6a8a" plt.rcParams.update({ "figure.facecolor": BG, "axes.facecolor": PANEL, "text.color": WHITE, "axes.labelcolor": WHITE, "xtick.color": DIM, "ytick.color": DIM, "axes.edgecolor": ACCENT, "grid.color": GRID_COL, "font.family": "DejaVu Sans", }) # ── ODE definition ────────────────────────────────────────────────────────── omega = 1.2 # rotation frequency lam = -0.25 # decay rate A = np.array([[0, -omega], [omega, 0 ]]) # pure rotation B = np.array([[lam, 0 ], [0, lam]]) # isotropic decay def vector_field(t, u): return (A + B) @ u # ── Matrix exponentials (closed-form for each step) ────────────────────── def exp_A(dt): """Exact matrix exponential of A*dt (rotation).""" c, s = np.cos(omega * dt), np.sin(omega * dt) return np.array([[c, -s], [s, c]]) def exp_B(dt): """Exact matrix exponential of B*dt (scaling).""" return np.exp(lam * dt) * np.eye(2) def exp_AB(dt): """Exact matrix exponential of (A+B)*dt.""" return expm((A + B) * dt) # ── Simulate trajectories ───────────────────────────────────────────────── u0 = np.array([1.5, 0.0]) T_total = 6.0 N_steps = 28 # splitting steps shown in animation dt = T_total / N_steps # Exact solution — dense ODE solve t_dense = np.linspace(0, T_total, 400) sol = solve_ivp(vector_field, [0, T_total], u0, t_eval=t_dense, method="RK45", rtol=1e-10, atol=1e-12) exact_x = sol.y[0] exact_y = sol.y[1] # Lie-Trotter: apply exp(A*dt) then exp(B*dt) at each step lt_points = [u0.copy()] u_lt = u0.copy() for _ in range(N_steps): u_lt = exp_A(dt) @ u_lt # step 1: rotate u_lt = exp_B(dt) @ u_lt # step 2: scale lt_points.append(u_lt.copy()) lt_points = np.array(lt_points) # Strang: exp(A*dt/2) * exp(B*dt) * exp(A*dt/2) st_points = [u0.copy()] u_st = u0.copy() for _ in range(N_steps): u_st = exp_A(dt/2) @ u_st # half-rotate u_st = exp_B(dt) @ u_st # full-scale u_st = exp_A(dt/2) @ u_st # half-rotate st_points.append(u_st.copy()) st_points = np.array(st_points) # Sub-steps for Lie-Trotter animation (show intermediate "elbow" point) lt_sub = [] # list of (after_A, after_B) per step u_lt2 = u0.copy() for _ in range(N_steps): after_A = exp_A(dt) @ u_lt2 after_B = exp_B(dt) @ after_A lt_sub.append((after_A.copy(), after_B.copy())) u_lt2 = after_B # Sub-steps for Strang (show three intermediate points) st_sub = [] u_st2 = u0.copy() for _ in range(N_steps): p1 = exp_A(dt/2) @ u_st2 p2 = exp_B(dt) @ p1 p3 = exp_A(dt/2) @ p2 st_sub.append((p1.copy(), p2.copy(), p3.copy())) u_st2 = p3 # ── Layout ────────────────────────────────────────────────────────────────── fig = plt.figure(figsize=(10.8, 10.8), dpi=100) fig.patch.set_facecolor(BG) # Main trajectory axis ax = fig.add_axes([0.08, 0.22, 0.84, 0.65]) ax.set_facecolor(PANEL) ax.set_aspect("equal") ax.set_xlim(-2.2, 2.2) ax.set_ylim(-2.2, 2.2) ax.grid(True, linewidth=0.4, alpha=0.5) ax.set_xlabel("x₁", fontsize=12, color=DIM) ax.set_ylabel("x₂", fontsize=12, color=DIM) # Title fig.text(0.5, 0.93, "Operator Splitting: Lie-Trotter vs Strang vs Exact", ha="center", va="center", fontsize=15, color=WHITE, fontweight="bold") fig.text(0.5, 0.895, r"$\dot{u} = (A + B)\,u$ · A = rotation, B = decay", ha="center", va="center", fontsize=11, color=DIM) # Legend legend_elements = [ Line2D([0], [0], color=EXACT, lw=2.5, label="Exact solution"), Line2D([0], [0], color=LT_COL, lw=2.0, linestyle="--", label="Lie-Trotter (1st order): eᴬᵈᵗ · eᴮᵈᵗ"), Line2D([0], [0], color=STRANG, lw=2.0, linestyle="-.", label="Strang splitting (2nd order): eᴬᵈᵗ/² · eᴮᵈᵗ · eᴬᵈᵗ/²"), ] ax.legend(handles=legend_elements, loc="upper right", fontsize=9, facecolor=ACCENT, edgecolor=GRID_COL, labelcolor=WHITE, framealpha=0.85) # Origin dot ax.plot(0, 0, "o", color=GRID_COL, markersize=3, zorder=1) # Draw full exact trajectory (background, faint) ax.plot(exact_x, exact_y, color=EXACT, lw=1.0, alpha=0.20, zorder=2) # Start point ax.plot(u0[0], u0[1], "o", color=WHITE, markersize=8, zorder=10) ax.annotate("start", xy=u0, xytext=(u0[0]+0.15, u0[1]+0.15), fontsize=9, color=WHITE, alpha=0.8) # ── Animated artists ───────────────────────────────────────────────────────── exact_line, = ax.plot([], [], color=EXACT, lw=2.5, zorder=5, alpha=0.95) lt_line, = ax.plot([], [], color=LT_COL, lw=1.8, zorder=6, linestyle="--", alpha=0.9) st_line, = ax.plot([], [], color=STRANG, lw=1.8, zorder=6, linestyle="-.", alpha=0.9) # Sub-step "elbow" segments shown during current frame lt_elbow_A, = ax.plot([], [], color=LT_COL, lw=1.2, zorder=7, alpha=0.6, linestyle=":") lt_elbow_B, = ax.plot([], [], color=LT_COL, lw=1.2, zorder=7, alpha=0.6, linestyle=":") st_sub1, = ax.plot([], [], color=STRANG, lw=1.2, zorder=7, alpha=0.6, linestyle=":") st_sub2, = ax.plot([], [], color=STRANG, lw=1.2, zorder=7, alpha=0.6, linestyle=":") st_sub3, = ax.plot([], [], color=STRANG, lw=1.2, zorder=7, alpha=0.6, linestyle=":") lt_dot, = ax.plot([], [], "o", color=LT_COL, markersize=8, zorder=9) st_dot, = ax.plot([], [], "o", color=STRANG, markersize=8, zorder=9) exact_dot, = ax.plot([], [], "o", color=EXACT, markersize=8, zorder=9) # Info text — bottom panel info_ax = fig.add_axes([0.0, 0.0, 1.0, 0.20]) info_ax.set_facecolor(BG) info_ax.axis("off") step_text = info_ax.text(0.5, 0.75, "", ha="center", va="center", fontsize=13, color=WHITE, fontweight="bold", transform=info_ax.transAxes) error_text = info_ax.text(0.5, 0.42, "", ha="center", va="center", fontsize=11, color=DIM, transform=info_ax.transAxes) method_text = info_ax.text(0.5, 0.10, "", ha="center", va="center", fontsize=10, color=DIM, style="italic", transform=info_ax.transAxes) # ── Helper: exact position at time t ──────────────────────────────────────── def exact_at(t): return expm((A + B) * t) @ u0 # ── Animation functions ────────────────────────────────────────────────────── def init(): for artist in [exact_line, lt_line, st_line, lt_elbow_A, lt_elbow_B, st_sub1, st_sub2, st_sub3, lt_dot, st_dot, exact_dot]: artist.set_data([], []) step_text.set_text("") error_text.set_text("") method_text.set_text("") return (exact_line, lt_line, st_line, lt_elbow_A, lt_elbow_B, st_sub1, st_sub2, st_sub3, lt_dot, st_dot, exact_dot, step_text, error_text, method_text) def update(frame): # frame = 0 is a "blank" hold; real frames start at 1 if frame == 0: return init() k = min(frame, N_steps) # current step index (capped) # ── Exact path up to t = k*dt ── t_now = k * dt t_sub = t_dense[t_dense <= t_now + 1e-12] x_sub = exact_x[:len(t_sub)] y_sub = exact_y[:len(t_sub)] exact_line.set_data(x_sub, y_sub) # ── Lie-Trotter path ── lt_line.set_data(lt_points[:k+1, 0], lt_points[:k+1, 1]) # ── Strang path ── st_line.set_data(st_points[:k+1, 0], st_points[:k+1, 1]) # ── Sub-step elbows for current step k (if not at the very start) ── if 1 <= k <= N_steps: idx = k - 1 prev_lt = lt_points[idx] mid_lt, end_lt = lt_sub[idx] # Elbow A: prev → mid_lt (rotate) lt_elbow_A.set_data([prev_lt[0], mid_lt[0]], [prev_lt[1], mid_lt[1]]) # Elbow B: mid_lt → end_lt (scale) — same as lt_line last segment, shown as dot lt_elbow_B.set_data([mid_lt[0], end_lt[0]], [mid_lt[1], end_lt[1]]) prev_st = st_points[idx] p1, p2, p3 = st_sub[idx] st_sub1.set_data([prev_st[0], p1[0]], [prev_st[1], p1[1]]) st_sub2.set_data([p1[0], p2[0]], [p1[1], p2[1]]) st_sub3.set_data([p2[0], p3[0]], [p2[1], p3[1]]) else: for seg in [lt_elbow_A, lt_elbow_B, st_sub1, st_sub2, st_sub3]: seg.set_data([], []) # ── Current dots ── ex_pos = exact_at(t_now) lt_dot.set_data([lt_points[k, 0]], [lt_points[k, 1]]) st_dot.set_data([st_points[k, 0]], [st_points[k, 1]]) exact_dot.set_data([ex_pos[0]], [ex_pos[1]]) # ── Errors ── err_lt = np.linalg.norm(lt_points[k] - ex_pos) err_st = np.linalg.norm(st_points[k] - ex_pos) step_text.set_text(f"Step {k} / {N_steps} · t = {t_now:.2f}") error_text.set_text( f"Error | Lie-Trotter: {err_lt:.4f} · Strang: {err_st:.4f}" ) # Cycle through "what's happening now" labels actions = [ "Lie-Trotter: rotate (A) → scale (B) | " "Strang: rotate½ → scale → rotate½", "Strang doubles order by symmetrising the steps", "Higher order = smaller error for the same Δt", ] method_text.set_text(actions[k % len(actions)]) return (exact_line, lt_line, st_line, lt_elbow_A, lt_elbow_B, st_sub1, st_sub2, st_sub3, lt_dot, st_dot, exact_dot, step_text, error_text, method_text) # ── Build & save animation ─────────────────────────────────────────────────── total_frames = N_steps + 3 # +3 for a hold at start/end ani = animation.FuncAnimation( fig, update, frames=total_frames, init_func=init, interval=500, # ms between frames (fps=2) blit=True, ) out_path = "/root/Strategy/content/drafts/splitting_animation.mp4" writer = animation.FFMpegWriter(fps=2, bitrate=2000, extra_args=["-vcodec", "libx264", "-pix_fmt", "yuv420p"]) ani.save(out_path, writer=writer, dpi=100) print(f"Saved: {out_path}") plt.close(fig)