Scientific Computing in Lean 4
Tomáš Skřivan
Carnegie Mellon University
Why Scientific Computing in Lean?
- scientific computing is full of mathematics + Lean understands mathematics
- Lean 4 is efficient general purpose programming language
- Lean metaprograming is very powerful and easy to use
- mathlib - a big library of formalized mathematics
- something interesting has to come out of this
SciLean
library for scientific computing in Lean 4
Motivation: Physics Simulation in Computer Graphics
work done at
Motivation
Sources of frustration:
- the mathematics is not in the code
- data often form algebraic structures
- composability and working with libraries
- libraries often impose data structures on the input output
- lack of well specified interfaces
-
prototyping
- mathematical description is often quite easy but implementation is not
- multi physics problems
- very difficult to make two solvers to talks to each other
- usually ends up in extremely complicated monolithic code
This video is from the SIGGRAPH 2022 Technical Paper: ‘Loki: a unified multiphysics simulation framework for production’.
Scientific Computing and Lean
What can Lean offer?
- the mathematics is not in the code
- the code is the math and the math is the code
- interactivity can effectively offer computer algebra system operating directly on the code
- composability and working with libraries
- interfaces can be given on mathematical level rather on the data level
- approximations compose badly
- interfaces be specified very precisely
-
prototyping
- interactively transform specification into code
- multi physics problems
- combining specifications is likely easier then combining implementations
Example: Harmonic Oscillator
x˙p˙=∂p∂H=−∂x∂H
\begin{aligned}
\dot x &= \frac{\partial H}{\partial p} \\
\dot p &= - \frac{\partial H}{\partial x}
\end{aligned}
def H (m k : ℝ) (x p : ℝ) := (1/(2*m)) * ∥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)))
(x, p) := solver m k substeps t (x, p) Δt
∇ (x’:=x), f x’=dx′dfx′=x
\texttt{∇ (x':=x), f x'} = \frac{d f}{d x'} \big|_{x'=x}
x(t)=odeSolve f t₀ x₀ t
⇔
x˙(t)=f(t,x(t)) x(t0)=x0
Two Main Operations of AD: Differential and Adjoint
Differential ∂:(X→Y)→(X→X→Y)
Adjoint †:(X→Y)→(Y→X)
∂fxdx=limh→0hf(x+h⋅dx)−f(x)
\partial \, f\, x \, dx = \lim_{h\rightarrow 0} \frac{f (x + h \cdot dx) - f(x)}{h}
x:X
f
f
y:Y
x:X
∂f
dy:Y
dx:X
∂
x:X
f
f
y:Y
y:Y
f
f†
x:X
†
∀xy,⟨f(x),y⟩=⟨x,f†(y)⟩
\forall \, x \, y, \left\langle f(x), y \right\rangle = \left\langle x, f^\dagger (y) \right\rangle
Simplifier and Typeclass approach
Differential ∂
Adjoint †
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) := ...
Other operators
- derivative f′ for f:R→X
- f′(x)=∂ f x 1
- gradient ∇f for f:X→R
- ∇f(x)=(∂ f x)† 1
- adjoint differential ∂†
- ∂† f x dy = (∂ f x)† dy
- tangent map T - forward mode AD
- T f x dx = (f x, ∂ f x dx)
- reverse differential R - reverse mode AD
- R f x = (f x, λ dy => (∂ f x)† dy)
- complexification fc
- fc:C→C for f:R→R
- inverse f−1
Forward Mode AD and Functoriality
g
∂
x:X
∂(f∘g)
dz:Z
dx:X
?
x:X
f
g
f
f
z:Z
∂f
x:X
dx:X
∂g
y:Y
dy:Y
dz:Z
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
T(f∘g)
dz:Z
dx:X
z:Z
x:X
Tg
dy:Y
dx:X
y:Y
Tf
dz:Z
z:Z
x:X
f
g
f
f
z:Z
Reverse Mode AD and Functoriality
x:X
f
g
?
f
f
z:Z
x:X
∂†(f∘g)
dx:X
dz:Z
∂†
x:X
∂†g
dx:X
y:Y
∂†f
dy:Y
dz:Z
Reverse Mode AD and Functoriality
x:X
f
g
?
f
f
z:Z
x:X
R(f∘g)
z:Z
∂†(f∘g)x:Z→X
∂†(f∘g)x:Z→X
ℛ (λ x => f (g x)) = λ x => let (y,dg') := ℛ g x let (z,df') := ℛ f y (z, λ dz => dg' (df' dz))
∂†
x:X
Rg
y:Y
∂†gx:Y→X
Rf
z:Z
∂†fy:Z→Y
R f x = (f x, λ dy => (∂ f x)† dy)
Example: Ballistic motion
x˙v˙x(0)=v=g−γf(v)=0v(0)=v0
\begin{aligned}
\dot x &= v \\
\dot v &= g - \gamma f(v) \\
x(0) &= 0 \qquad v(0) = v_0
\end{aligned}
For T∈R,xT∈R2 find v0 such that x(T)=xT
We will solve necessary condition
v0=argminv0∗∥x(T)−xT∥2where x˙(0)=v0∗
Problem:
To compute ∂v0∂x(T) we have to solve adjoint problem
x˙∗v˙∗xT∗(T)=0=x∗−γf′(v)v∗=x(T)−xTv∗(T)=0
\begin{aligned}
\dot x^* &= 0 \\
\dot v^* &= x^* - \gamma f'(v) v^* \\
x^*_T(T) &= x(T) - x_T \qquad v^*(T) = 0
\end{aligned}
∂v0∂x(T)=2v∗(0)
\frac{\partial x(T)}{\partial v_0} = 2 v^*(0)
Example: Ballistic motion
x˙v˙x(0)=v=g−γf(v)=0v(0)=v0
\begin{aligned}
\dot x &= v \\
\dot v &= g - \gamma f(v) \\
x(0) &= 0 \qquad v(0) = v_0
\end{aligned}
For T∈R,xT∈R2 find v0 such that x(T)=xT
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
def ballisticMotion (x v : ℝ×ℝ) := (v, g - (5 + ‖v‖) • v)
SciLean summary
- library for scientific computing
- aim is to help users with the mathematics
- provide guarantees when possible
- current focus is on automatic differentiation
- priorities:
- user experience
- performance
- correctness
- formal correctness
Scientific Computing in Lean 4 Tomáš Skřivan Carnegie Mellon University
Talk at Certified and Symbolic-Numeric Computation
By lecopivo
Talk at Certified and Symbolic-Numeric Computation
- 408
