Ecosyste.ms: Awesome

An open API service indexing awesome lists of open source software.

Awesome Lists | Featured Topics | Projects

https://github.com/lecopivo/SciLean

Scientific computing in Lean 4
https://github.com/lecopivo/SciLean

lean4 numerical-methods scientific-computing symbolic-computation

Last synced: 2 months ago
JSON representation

Scientific computing in Lean 4

Awesome Lists containing this project

README 

        

* SciLean: Scientific Computing Assistant



  Framework for scientific computing such as solving differential equations, optimization or machine learning written in [[http://leanprover.github.io/][Lean]]. This library is in an *extremely early stage of development* and at its current stage is just a proof of concept on how Lean can be used for scientific computing.



Lean is an expressive functional programming language that allows to formalize the mathematics behind these computations. This can offer several benefits:



    - Code transformation and optimization guided by formalization of underlining mathematics, like automatic differentiation, algebraic simplification, fine control of used approximations or execution scheduling.

    - First class symbolic computation. Any function can be purely symbolic, functions like `gradient`, `integral` or `limit` are inherently non-computable. However, they carry meaning what the program should be doing and we provide tools to manipulate them or approximate them with actually computable function.

    - Code generation based on formal specification. Many problems any scientific computing or machine learning can be stated very easily e.g. find a minimizer of a function. We then provide tools how to turn such specification into a runnable code satisfying the specification, usually in an appropriate limit of used approximations.

    - Catalogization of numerical methods.



  In short, mathematics is the ultimate abstraction for numerical computing and Lean can understand mathematics. Hopefully, using Lean will allow us to create really powerful and extensible library for scientific computing.



  

* Installation and running examples/tests



As we are using Lean programming language, you need Lean's version manager =elan=. Follow its installation [[https://github.com/leanprover/elan#installation][instructions]].



Getting and building SciLean simply:

#+begin_src

git clone https://github.com/lecopivo/SciLean.git

cd SciLean

lake exe cache get

lake build

#+end_src



To run examples:

#+begin_src

lake build HarmonicOscillator && .lake/build/bin/HarmonicOscillator

lake build WaveEquation && .lake/build/bin/WaveEquation 

#+end_src

Other examples in =examples= directory do not currently work.



* Example: Simulation of harmonic oscillator



  The idea behind this library is can be most easily demonstrated on an example. We will create simulation of a harmonic oscillator, full [[https://github.com/lecopivo/SciLean/blob/master/examples/HarmonicOscillator.lean][example]].



  Harmonic oscillator is defined by its energy/Hamiltonian. Lean allows us to write down the energy in very familiar form:

#+begin_src

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

#+end_src



  One of the main novelties of SciLean is that it allows us to formally state our goal i.e. simulation of harmonic oscillator:

#+begin_src

def solver (m k : ℝ) (steps : Nat)

  : Impl (ode_solve (HamiltonianSystem (H m k))) := ...

#+end_src

  The above code declares a function =solver= taking mass =m=, spring stiffness =k= and number of steps =steps=. The type of the function =solver= is an implementation of ODE solver of Hamiltonian system given by Hamiltonian =(H m k)=. This means that it is a function taking =(t, x, p)= and producing state =(x, p)= by approximately evolving the specified Hamiltonian system. Furthermore, it has compiler checked guarantee that as you increase the number of =steps= the function output converges to the correct answer(not completely true but good enough explanation for now).



  The last three dots =...= is the place where the actual definition of the =solver= goes. This is a place where we completely diverge from traditional programming languages and where the word ~assistant~ from the title comes in. We do not tell computer what to do step by step but we guide it on how to transform the formal specification into an executable code. This sounds way to abstract, here is the actual code:

#+begin_src

def solver (m k : ℝ) (steps : Nat)

: Impl (ode_solve (HamiltonianSystem (H m k))) :=

by

-- Unfold Hamiltonian definition and compute gradients

simp[HamiltonianSystem, H];

autograd

autograd



-- Apply RK4 method

rw [ode_solve_fixed_dt runge_kutta4_step]

lift_limit steps "Number of ODE solver steps."; admit; simp



finish_impl

#+end_src



  Let's start with a warning! This code is not meant to be consumed as it is but it should be read interactively in a text editor, see the following animation:



  [[file:anim.gif]]



  On the right side you can see what is our current goal based on the position of the cursor. At the beginning, the goal is =Impl (ode_solve (HamiltonianSystem (H m k)))= as we have defined =solver= function.



  In the next step, we expand the definition of =HamiltonianSystem= and =H=. This allows us to proceed to symbolic differentiation step.



  Let's pause for a bit. If we expand only =HamiltonianSystem=, the goal would contain:

#+begin_src

fun x p => (gradient (H m k x) p, -gradient (fun x => H m k x p) x))

#+end_src

  These are exactly [[https://en.wikipedia.org/wiki/Hamiltonian_mechanics#From_Euler-Lagrange_equation_to_Hamilton's_equations][Hamilton's equations of motion]]. The first term =gradient (H m k x) p= is derivative of Hamiltonian w.r.t. to momentum and the second term =gradient (fun x => H m k x p) x= it derivative of Hamiltonian w.r.t. to position.



  Thus after expanding =H= as well we call =autograd= twice to eliminate both =gradient= functions.



  The =gradient= function is purely symbolic. A term like =gradient f= just tells us we want to have a program that computes the gradient of a function =f=. Our task in the interactive mode is to eliminate every occurrence of =gradient f= (or any other purely symbolic function) and replace it with an actually computable function. This can be done through symbolic differentiation, automatic differentiation(like forward or backward mode differentiation) or approximate it with finite differences.



  The second purely symbolic function in our specification is =ode_solve=. We have to decide which ODE integration scheme we want to use. Let's pick fixed time step Runge-Kutta 4:  

#+begin_src

rw [ode_solve_fixed_dt runge_kutta4_step]

#+end_src



  At this point the goal is still saying that we are implementing the original goal =Impl (ode_solve (HamiltonianSystem (H m k)))= in a limit of infinite number of ODE integrator steps. The next line just accepts the fact that we can not compute the answer exactly and pick a concrete number of steps:

#+begin_src

lift_limit steps "Number of ODE solver steps."; admit; simp

#+end_src



The last line =finish_impl= states we are done. Lean checks if all purely symbolic function have been eliminated and the actually executable code can be generated.



See the full [[https://github.com/lecopivo/SciLean/blob/master/examples/HarmonicOscillator.lean][example]] to see how the set up initial conditions and how the function =solver= is actually used. To execute this example, run:

#+begin_src

lake env lean --run examples/HarmonicOscillator.lean 

#+end_src

from the project root directory.



The simulation of harmonic oscillator is not very interesting. The great thing about SciLean is that we can simply change the definition of Hamiltonian and get simulation of way more interesting systems. The Hamiltonian for wave equation(discretized in space) is:

#+begin_src

def H (m k : ℝ) (x p : ℝ^n) := 

  let Δx := (1 : ℝ)/(n : ℝ)

  (Δx/(2*m)) * ∥p∥² + (Δx * k/2) * (∑ i, ∥x[i] - x[i-1]∥²)

#+end_src

The rest of the code can stay the same. Running this [[https://github.com/lecopivo/SciLean/blob/master/examples/WaveEquation.lean][example]]:

#+begin_src

lake env lean --run examples/WaveEquation.lean

#+end_src

Produces the following animation:

[[file:wave.gif]]



Or when nicely rendered:



[[file:wavering.gif]]