Improve MAM tutorial: add problem statement, restructure continuation, use @init macro

Author: ocots

docs/src/tutorial-mam.md

@@ -1,7 +1,12 @@
 # Minimal Action Method using Optimal Control
 
-The Minimal Action Method is a numerical technique for finding the most probable transition pathway between stable states in stochastic dynamical systems. It achieves this by minimizing an action functional that represents the path's deviation from the deterministic dynamics, effectively identifying the path of least resistance through the system's landscape.
-This tutorial demonstrates how to implement MAM as an optimal control problem.
+```@meta
+Draft = false
+```
+
+The Minimal Action Method (MAM) is a numerical technique for finding the most probable transition pathway between stable states in stochastic dynamical systems. It achieves this by minimizing an action functional that represents the path's deviation from the deterministic dynamics, effectively identifying the path of least resistance through the system's landscape.
+
+This tutorial demonstrates how to implement MAM as an optimal control problem, using the classical Maier-Stein model as a benchmark example.
 
 ## Required Packages
 
@@ -11,9 +16,29 @@ using NLPModelsIpopt
 using Plots, Printf
 ```
 
+## Problem Statement
+
+We aim to find the most probable transition path between two stable states of a stochastic dynamical system. For a system with deterministic dynamics $f(x)$ and small noise, the transition path minimizes the action functional:
+
+```math
+S[x(\cdot), u(\cdot)] = \int_0^T \|u(t) - f(x(t))\|^2 \, dt
+```
+
+subject to the path dynamics:
+
+```math
+\dot{x}(t) = u(t), \quad x(0) = x_0, \quad x(T) = x_f
+```
+
+where $x_0$ and $x_f$ are the initial and final states, and $T$ is the transition time.
+
+!!! note "Physical interpretation"
+
+    The action $S$ measures the "cost" of deviating from the deterministic flow $f(x)$. Paths with smaller action are exponentially more likely in the small noise limit.
+
 ## Problem Setup
 
-We'll consider a 2D system with a double-well flow, called the Maier-Stein model. It is a famous benchmark problem as it exhibits non-gradient dynamics with two stable equilibrium points at (-1,0) and (1,0), connected by a non-trivial transition path.
+We consider a 2D system with a double-well flow, called the Maier-Stein model. It is a famous benchmark problem as it exhibits non-gradient dynamics with two stable equilibrium points at $(-1,0)$ and $(1,0)$, connected by a non-trivial transition path.
 The system's deterministic dynamics are given by:
 
 ```@example main-mam
@@ -45,29 +70,39 @@ nothing # hide
 
 ## Initial Guess
 
-We provide an initial guess for the path using a simple interpolation:
+We provide an initial guess for the path using a simple interpolation with the `@init` macro:
 
 ```@example main-mam
 # Time horizon
 T = 50
 
-# Linear interpolation for x₁
-x1(t) = -(1 - t/T) + t/T
-
-# Parabolic guess for x₂
-x2(t) = 0.3(-x1(t)^2 + 1)
-x(t) = [x1(t), x2(t)]
-u(t) = f(x(t))
-
-# Initial guess
-init = (state=x, control=u)
+# Helper functions for initial state guess
+L(t) = -(1 - t/T) + t/T      # Linear interpolation from -1 to 1
+P(t) = 0.3*(-L(t)^2 + 1)     # Parabolic arc (approximates saddle crossing)
+
+init = @init ocp(T) begin
+    # Linear interpolation for x₁
+    x₁(t) := L(t)
+    # Parabolic guess for x₂
+    x₂(t) := P(t)
+    # Control from vector field
+    u(t) := f(L(t), P(t))
+end
 nothing # hide
 ```
 
+!!! note "Initial guess strategy"
+
+    The initial guess uses a simple geometric path: linear interpolation in $x_1$ and a parabolic arc in $x_2$. This provides a reasonable starting point that avoids the unstable saddle point at the origin. The control is initialized to follow the deterministic flow along this path.
+
 ## Solving the Problem
 
 We solve the problem in two steps for better accuracy:
 
+!!! note "Two-step resolution"
+
+    Starting with a coarse grid (50 points) allows for faster initial convergence. Refining with a fine grid (1000 points) then improves accuracy of the solution.
+
 ```@example main-mam
 # First solve with coarse grid
 sol = solve(ocp(T); init=init, grid_size=50)
@@ -104,22 +139,50 @@ The resulting path shows the most likely transition between the two stable state
 To find the maximum likelihood path, we also need to minimize the transient time `T`. Hence, we perform a discrete continuation over the parameter `T` by solving the optimal control problem over a continuous range of final times `T`, using each solution to initialize the next problem.
 
 ```@example main-mam
-objectives = []
-Ts = range(1,100,100)
-sol = solve(ocp(Ts[1]); display=false, init=init, grid_size=200)
-println(" Time   Objective     Iterations")
-for T=Ts
-    global sol = solve(ocp(T); display=false, init=sol, grid_size=1000, tol=1e-8)
-    @printf("%6.2f  %9.6e  %d\n", T, objective(sol), iterations(sol))
-    push!(objectives, objective(sol))
+# Continuation function to avoid global variables
+function continuation_mam(Ts; init_guess=init)
+    objectives = Float64[]
+    current_sol = init_guess
+    
+    println(" Time   Objective     Iterations")
+    for T in Ts
+        current_sol = solve(ocp(T); display=false, init=current_sol, grid_size=1000, tol=1e-8)
+        obj = objective(current_sol)
+        @printf("%6.2f  %9.6e  %d\n", T, obj, iterations(current_sol))
+        push!(objectives, obj)
+    end
+    
+    return objectives, current_sol
 end
+
+# Perform continuation
+Ts = range(1, 100, 100)
+objectives, final_sol = continuation_mam(Ts)
+nothing # hide
 ```
 
+We can now analyze the results and identify the optimal transition time:
+
+```@example main-mam
+# Find optimal time
+idx_min = argmin(objectives)
+T_min = Ts[idx_min]
+obj_min = objectives[idx_min]
+
+@printf("Optimal transition time: T* = %.2f\n", T_min)
+@printf("Minimal action: S* = %.6e\n", obj_min)
+```
+
+Let us visualize the evolution of the objective function with respect to the transition time:
+
 ```@example main-mam
-T_min = Ts[argmin(objectives)]
 plt1 = scatter(Ts, log10.(objectives), xlabel="Time", label="Objective (log10)")
-vline!(plt1, [T_min], label="Minimum", z_order=:back)
+vline!(plt1, [T_min], label="Minimum at T=$(round(T_min, digits=1))", z_order=:back)
 plt2 = scatter(Ts[40:100], log10.(objectives[40:100]), xlabel="Time", label="Objective (log10)")
 vline!(plt2, [T_min], label="Minimum", z_order=:back)
 plot(plt1, plt2, layout=(2,1), size=(800,800))
 ```
+
+!!! note "Interpretation"
+
+    The optimal transition time $T^*$ balances two competing effects: shorter times require larger deviations from the deterministic flow (higher action), while longer times allow the system to follow the flow more closely. The minimum represents the most probable transition time in the small noise limit.