SciLean: Scientific Computing Assistant

Framework for scientific computing such as solving differential equations, optimization or machine learning written in 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 powerfull 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 instructions.

Getting and building SciLean simply:

git clone https://github.com/lecopivo/SciLean.git
cd SciLean
lake build

To run examples:

lake env lean --run examples/HarmonicOscilator.lean 
lake env lean --run examples/WaveEquation.lean 

Other examples in examples directory do not currently work.

To run tests:

lake run tests

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 example.

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

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

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

def solver (m k : ℝ) (steps : Nat)
  : Impl (ode_solve (HamiltonianSystem (H m k))) := ...

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:

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

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:

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:

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

These are exactly 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:

rw [ode_solve_fixed_dt runge_kutta4_step]

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:

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

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 example to see how the set up initial conditions and how the function solver is actually used. To execute this example, run:

lake env lean --run examples/HarmonicOscilator.lean 

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:

def H (m k : ℝ) (x p : ℝ^n) := 
  let Δx := (1 : ℝ)/(n : ℝ)
  (Δx/(2*m)) * ∥p∥² + (Δx * k/2) * (∑ i, ∥x[i] - x[i-1]∥²)

The rest of the code can stay the same. Running this example:

lake env lean --run examples/WaveEquation.lean

Produces the following animation:

Or when nicely rendered: