Talk at Certified and Symbolic-Numeric Computation

Scientific Computing in Lean 4

Tomáš Skřivan

Carnegie Mellon University

Other operators

derivative $f'$ for $f : \mathbb{R} \rightarrow X$
- $f'(x) = \texttt{∂ f x 1}$
gradient $\nabla f$ for $f : X \rightarrow \mathbb{R}$
- $\nabla f(x) = \texttt{(∂ f x)† 1}$
adjoint differential $\partial^\dagger$
- $\texttt{∂† f x dy = (∂ f x)† dy}$
tangent map $\mathcal{T}$ - forward mode AD
- $\mathcal{T}\texttt{ f x dx = (f x, ∂ f x dx)}$
reverse differential $\mathcal{R}$ - reverse mode AD
- $\mathcal{R}\texttt{ f x = (f x, λ dy => (∂ f x)† dy)}$
complexification $f^c$
- $f^c : \mathbb{C} \rightarrow \mathbb{C}$ for $f : \mathbb{R} \rightarrow \mathbb{R}$
inverse $f^{-1}$

Forward Mode AD and Functoriality

 𝒯 (λ x => f (g x)) 
= 
λ x dx =>
  let (y,dy) := 𝒯 g x dx
  let (z,dz) := 𝒯 f y dy
  (z,dz)

$x:X$

$\mathcal{T} (f\circ g)$

$dz:Z$

$dx:X$

$z:Z$

$x:X$

$\mathcal{T} g$

$dy:Y$

$dx:X$

$y:Y$

$\mathcal{T} f$

$dz:Z$

$z:Z$

$x:X$

$f$

$g$

$f$

$z:Z$

Reverse Mode AD and Functoriality

$x:X$

$f$

$g$

$?$

$f$

$z:Z$

$x:X$

$\mathcal{R} (f\circ g)$

$z:Z$

$\partial^\dagger (f\circ g) \, x : Z \rightarrow X$

 ℛ (λ x => f (g x)) 
    = 
    λ x => 
      let (y,dg') := ℛ g x
      let (z,df') := ℛ f y
      (z, λ dz => dg' (df' dz))

$\partial^\dagger$

$x:X$

$\mathcal{R} g$

$y:Y$

$\partial^\dagger g \, x : Y \rightarrow X$

$\mathcal{R} f$

$z:Z$

$\partial^\dagger f \, y : Z \rightarrow Y$

$\mathcal{R}\texttt{ f x = (f x, λ dy => (∂ f x)† dy)}$

Scientific Computing in Lean 4 Tomáš Skřivan Carnegie Mellon University

	def H (m k : ℝ) (x p : ℝ) := (1/(2m)) ∥p∥² + k/2 * ∥x∥²

	approx solver (m k : ℝ) (steps : Nat)
	:= odeSolve (λ t (x,p) => ( ∇ (p':=p), H m k x p',
	-∇ (x':=x), H m k x' p))
	by
	-- Unfold Hamiltonian and compute gradients
	unfold H
	symdiff; symdiff

	-- Apply RK4 method
	rw [odeSolve_fixed_dt runge_kutta4_step]
	approx_limit steps; simp; intro steps';

	Goals (1)
	m k : ℝ
	steps : ℕ
	⊢ Approx (odeSolve fun t x => (1 / m * x.snd, -(k * x.fst)))

	Goals (1)
	m k : ℝ
	steps : ℕ
	⊢ Approx (odeSolve (λ t (x,p) => ( ∇ (p':=p), H m k x p',
	-∇ (x':=x), H m k x' p)))

	Goals (1)
	m k : ℝ
	steps steps' : ℕ
	⊢ Approx (odeSolve_fixed_dt_impl steps' runge_kutta4_step
	fun t x => (1 / m * x.snd, -(k * x.fst)))

	instance comp_is_smooth (f : Y → Z) (g : X → Y) [IsSmooth f] [IsSmooth g]
	: IsSmooth (λ x => f (g x)) := ...

	@[simp]
	theorem chain_rule (f : Y → Z) (g : X → Y) [IsSmooth f] [IsSmooth g]
	: ∂ (λ x => f (g x))
	=
	λ x dx => ∂ f (g x) (∂ g x dx) := ...

	instance comp_has_adjoint (f : Y → Z) (g : X → Y) [IsLin f] [IsLin g]
	: IsLin (λ x => f (g x)) := ...

	@[simp]
	theorem adj_of_comp (f : Y → Z) (g : X → Y) [IsLin f] [IsLin g]
	: (λ x => f (g x))†
	=
	λ z => g† (f† z) := ...

	𝒯 (λ x => f (g x))
	=
	λ x dx =>
	let (y,dy) := 𝒯 g x dx
	let (z,dz) := 𝒯 f y dy
	(z,dz)

	ℛ (λ x => f (g x))
	=
	λ x =>
	let (y,dg') := ℛ g x
	let (z,df') := ℛ f y
	(z, λ dz => dg' (df' dz))

	approx aimToTarget (init_v : ℝ×ℝ) (optimizationRate : ℝ) :=
	λ (T : ℝ) (target : ℝ×ℝ) =>
	let shoot := λ (t : ℝ) (v : ℝ×ℝ) =>
	odeSolve (t₀ := 0) (x₀ := ((0:ℝ×ℝ),v)) (t := t)
	(f := λ (t : ℝ) (x,v) => balisticMotion x v)
	(λ v => (shoot T v).1)⁻¹ target