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

\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
\texttt{∇ (x':=x), f x'} = \frac{d f}{d x'} \big|_{x'=x}

\(x(t) = \texttt{odeSolve f t₀ x₀ t}\)

\(\Leftrightarrow\)

\(\dot x(t) = f(t,x(t))     x(t_0) = x_0\)

Two Main Operations of AD: Differential and Adjoint

Differential \(\partial : (X \rightarrow Y) \rightarrow (X \rightarrow X \rightarrow Y)\)

Adjoint \(\dagger : (X \rightarrow Y) \rightarrow (Y \rightarrow 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\)

\(\partial f \)

\(dy:Y\)

\(dx:X\)

 \(\partial\)

\(x:X\)

\(f \)

\(f \)

\(y:Y\)

\(y:Y\)

\(f \)

\(f^\dagger \)

\(x:X\)

\(\dagger\)

\forall \, x \, y, \left\langle f(x), y \right\rangle = \left\langle x, f^\dagger (y) \right\rangle

Simplifier and Typeclass approach

Differential \(\partial\)

Adjoint \(\dagger\)

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 : \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

\(g \)

 \(\partial\)

\(x:X\)

\(\partial (f\circ g) \)

\(dz:Z\)

\(dx:X\)

 \(?\)

\(x:X\)

\(f \)

\(g \)

\(f \)

\(f \)

\(z:Z\)

\(\partial f \)

\(x:X\)

\(dx:X\)

\(\partial 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\)

\(\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 \)

\(f \)

\(z:Z\)

Reverse Mode AD and Functoriality

\(x:X\)

\(f \)

\(g \)

 \(?\)

\(f \)

\(f \)

\(z:Z\)

\(x:X\)

\(\partial^\dagger (f\circ g) \)

\(dx:X\)

\(dz:Z\)

 \(\partial^\dagger\)

\(x:X\)

\(\partial^\dagger g \)

\(dx:X\)

\(y:Y\)

\(\partial^\dagger f \)

\(dy:Y\)

\(dz:Z\)

Reverse Mode AD and Functoriality

\(x:X\)

\(f \)

\(g \)

 \(?\)

\(f \)

\(f \)

\(z:Z\)

\(x:X\)

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

\(z:Z\)

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

\(\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)}\)

Example: Ballistic motion

\begin{aligned} \dot x &= v \\ \dot v &= g - \gamma f(v) \\ x(0) &= 0 \qquad v(0) = v_0 \end{aligned}

For \(T\in \mathbb{R}, x_T \in \mathbb{R}^2\) find \(v_0\) such that \(x(T) = x_T\)

We will solve necessary condition

   \(v_0 = \argmin_{v_0^*} \|x(T) - x_T\|^2 \qquad \text{where } \dot x(0) = v_0^*\)

Problem:

To compute \(\frac{\partial x(T)}{\partial v_0}\) we have to solve adjoint problem

\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}
\frac{\partial x(T)}{\partial v_0} = 2 v^*(0)

Example: Ballistic motion

\begin{aligned} \dot x &= v \\ \dot v &= g - \gamma f(v) \\ x(0) &= 0 \qquad v(0) = v_0 \end{aligned}

For \(T\in \mathbb{R}, x_T \in \mathbb{R}^2\) find \(v_0\) such that \(x(T) = x_T\)

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:
    1. user experience
    2. performance
    3. correctness
    4. formal correctness