dimensional analysis and arithmetic in type

Stream: lean4

[Tomas Skrivan]

There is a thread in #maths Dimensional Analysis Community Input. It prompted me to give a try at dimensional analysis in Lean 4.

I have this working:

def motion (a : ℝ[m*s⁻²]) (t : ℝ[s]) : ℝ[m] := (1/2:ℝ) * a * t * t

but I have trouble with addition, this does not work

def motion' (v : ℝ[m*s⁻¹]) (a : ℝ[m*s⁻²]) (t : ℝ[s]) : ℝ[m] := (1/2: ℝ) * a * t * t + v * t

I'm getting an error that instance of HAdd can not be found:

/home/tomass/doodle/lean/units.lean:94:63: error: failed to synthesize instance
  HAdd
    (Dimensional ℝ
      ({ s := 0 + -2 * 1, m := 1 + -2 * 0, kg := 0 + -2 * 0 } * { s := 1, m := 0, kg := 0 } *
        { s := 1, m := 0, kg := 0 }))
    (Dimensional ℝ ({ s := 0 + -1, m := 1 + -0, kg := 0 + -0 } * { s := 1, m := 0, kg := 0 })) ?m.10569

You can clearly see that the arithmetic stayed in the type and it is tripping typeclass synthesis.

How can I deal with this problem?

[Tomas Skrivan]

The code:

structure SIUnit where
  s := (0 : Int)
  m := (0 : Int)
  kg := (0 : Int)
deriving DecidableEq

def meter : SIUnit := { m := 1}
def second : SIUnit := { s := 1 }
def kilogram : SIUnit := { kg := 1 }

instance : Mul SIUnit :=
  ⟨λ x y => { s := x.s + y.s,
              m := x.m + y.m,
              kg := x.kg + y.kg }⟩

instance : HPow SIUnit Int SIUnit :=
  ⟨λ x y => { s := x.s * y,
              m := x.m * y,
              kg := x.kg * y }⟩

instance : OfNat SIUnit 1 := ⟨{}⟩

structure Dimensional (α : Type u) (units : SIUnit) where
  val : α

instance {α units} [Add α] : Add (Dimensional α units) := ⟨λ x y => ⟨x.val + y.val⟩⟩
instance {α units} [OfNat α n] : OfNat (Dimensional α units) n := ⟨⟨OfNat.ofNat n⟩⟩

instance {α β γ units units'} [HMul α β γ]
  : HMul (Dimensional α units) (Dimensional β units') (Dimensional γ (units*units')) :=
  ⟨λ x y => ⟨x.val * y.val⟩⟩
instance {α β γ units} [HMul α β γ]
  : HMul (Dimensional α units) β (Dimensional γ (units)) :=
  ⟨λ x y => ⟨x.val * y⟩⟩
instance {α β γ units'} [HMul α β γ]
  : HMul α (Dimensional β units') (Dimensional γ (units')) :=
  ⟨λ x y => ⟨x * y.val⟩⟩

def Dimensional.cast {α units units'} (a : Dimensional α units) (h : units = units')
  : Dimensional α units' := h ▸ a

abbrev ℝ := Float

declare_syntax_cat siunit (behavior := both)

syntax "m" : siunit
syntax "s" : siunit
syntax "kg" : siunit
syntax siunit "*" siunit : siunit
syntax siunit noWs "²" : siunit
syntax siunit noWs "⁻¹" : siunit
syntax siunit noWs "⁻²" : siunit

syntax "{" term "," term "," term "}"  : siunit

syntax "ℝ[" siunit "]" : term

macro_rules
| `(siunit| m) => `(siunit| {1,0,0} )
| `(siunit| s) => `(siunit| {0,1,0} )
| `(siunit| kg) => `(siunit| {0,0,1} )
| `(siunit| $unit * $unit') => do
  match (← Lean.expandMacros unit) with
  | `(siunit| {$m, $s, $kg}) =>
    match (← Lean.expandMacros unit') with
    | `(siunit| {$m', $s', $kg'}) =>
      `(siunit| {$m + $m', $s + $s', $kg + $kg'})
    | _ => Lean.Macro.throwUnsupported
  | _ => Lean.Macro.throwUnsupported
| `(siunit| $unit²) => do
  match (← Lean.expandMacros unit) with
  | `(siunit| {$m, $s, $kg}) =>
    `(siunit| {2*$m, 2*$s, 2*$kg})
  | _ => Lean.Macro.throwUnsupported
| `(siunit| $unit⁻¹) => do
  match (← Lean.expandMacros unit) with
  | `(siunit| {$m, $s, $kg}) =>
    `(siunit| {-$m, -$s, -$kg})
  | _ => Lean.Macro.throwUnsupported
| `(siunit| $unit⁻²) => do
  match (← Lean.expandMacros unit) with
  | `(siunit| {$m, $s, $kg}) =>
    `(siunit| {-2*$m, -2*$s, -2*$kg})
  | _ => Lean.Macro.throwUnsupported
| `(ℝ[ $units ]) => do
  match (← Lean.expandMacros units) with
  | `(siunit| {$m, $s, $kg}) => `(Dimensional ℝ {m:=$m, s:=$s, kg:=$kg})
  | _ => Lean.Macro.throwUnsupported


def motion (a : ℝ[m*s⁻²]) (t : ℝ[s]) : ℝ[m] := (1/2:ℝ) * a * t * t


def motion' (v : ℝ[m*s⁻¹]) (a : ℝ[m*s⁻²]) (t : ℝ[s]) : ℝ[m] := (1/2: ℝ) * a * t * t + v * t

-- this works but is ugly
def motion'' (v : ℝ[m*s⁻¹]) (a : ℝ[m*s⁻²]) (t : ℝ[s]) : ℝ[m] := (1/2: ℝ) * a * t * t + (v * t).cast (by decide)

[Kyle Miller]

One thing you can to to ease the pain a bit is

def Dimensional.cast {α units units'} (a : Dimensional α units)
  (h : units = units' := by (first | decide | simp))
  : Dimensional α units' := h ▸ a

since then you can write

def motion'' (v : ℝ[m*s⁻¹]) (a : ℝ[m*s⁻²]) (t : ℝ[s]) : ℝ[m] := (1/2: ℝ) * a * t * t + (v * t).cast

[Kyle Miller]

I threw in simp just to try to make it more likely to succeed.

[pcpthm]

A solution is to write a custom elaborator.

open Lean Lean.Meta Lean.Elab Lean.Elab.Term in
elab_rules : term
| `(binop% HAdd.hAdd $x $y) => do
  let x ← elabTerm x none
  let ty ← instantiateMVars (← inferType x)
  if ty.isAppOf ``Dimensional then
    let inst ← synthInstance (← mkAppM ``Add #[ty])
    let inst ← mkAppOptM ``instHAdd #[ty, inst]
    let y ← elabTerm y ty
    mkAppOptM ``HAdd.hAdd #[ty, ty, ty, inst, x, y]
  else
    throwUnsupportedSyntax

[Tomas Skrivan]

This is great!

This can probably be extended to solve:

1. for v : ℝ[m*s⁻¹] and t : ℝ[s] have v*t : ℝ[m] instead of Dimensional ℝ ({ s := 0 + -1, m := 1 + -0, kg := 0 + -0 } * { s := 1, m := 0, kg := 0 })
2. for t : ℝ[s] and f : ℝ[s⁻¹] have t*f : ℝ instead of Dimensional ℝ { s := 0, m := 0, kg := 0 }

[Tomas Skrivan]

Half way there, I have 1. but the resulting type is Dimensional ℝ { s := Int.ofNat 0, m := Int.ofNat 1, kg := Int.ofNat 0 }. How can I get 0 instead of Int.ofNat 0?

[Tomas Skrivan]

#check λ (v : ℝ[m*s⁻¹]) (t : ℝ[s]) => v*t

@[simp]
theorem SIUnit.mul_simp (m m' s s' kg kg')
  : SIUnit.mk m s kg * SIUnit.mk m' s' kg' = ⟨m+m', s+s', kg+kg'⟩ := by rfl

open Lean Lean.Meta Lean.Elab Lean.Elab.Term in
elab_rules : term
| `(binop% HMul.hMul $x $y) => do
  let x ← elabTerm x none
  let y ← elabTerm y none
  let tx ← instantiateMVars (← inferType x) -- type of x
  let ty ← instantiateMVars (← inferType y) -- type of y
  if tx.isAppOf ``Dimensional ∧ ty.isAppOf ``Dimensional then
    let z ← mkAppM ``HMul.hMul #[x,y]
    let tz ← instantiateMVars (← inferType z)
    let uz := tz.getArg! 1 -- units of z
    let uz ← reduce (← simp uz (← Simp.Context.mkDefault)).expr -- simplify units
    let tz ← mkAppOptM ``Dimensional #[tz.getArg! 0, uz]
    mkAppOptM ``HMul.hMul #[tx, ty, tz, none, x, y]
  else
    throwUnsupportedSyntax

#check λ (v : ℝ[m*s⁻¹]) (t : ℝ[s]) => v*t

[Tomas Skrivan]

Another useful elaborator for ℝ[m*s⁻¹]

#check ℝ[m*s⁻¹] -- Dimensional ℝ { s := 0 + -1, m := 1 + -0, kg := 0 + -0 }
#check ℝ[s⁻¹*m] -- Dimensional ℝ { s := -1 + 0, m := -0 + 1, kg := -0 + 0 }

open Lean Lean.Meta Lean.Elab Lean.Elab.Term in
elab_rules : term
| `(Dimensional ℝ {m:=$m, s:=$s, kg:=$kg}) => do
  let T ← reduce (← elabTerm (← `({m:=$m, s:=$s, kg:=$kg})) (mkConst ``SIUnit))
  mkAppM ``Dimensional #[mkConst ``ℝ, T]

#check ℝ[m*s⁻¹] -- Dimensional ℝ { s := Int.negSucc 0, m := Int.ofNat 1, kg := Int.ofNat 0 }
#check ℝ[s⁻¹*m] -- Dimensional ℝ { s := Int.negSucc 0, m := Int.ofNat 1, kg := Int.ofNat 0 }

[Kevin Buzzard]

These kinds of answers make me realise that I really still don't have a good understanding at all of the new possibilities which Lean 4 will give us.

[Tomas Skrivan]

And I'm constantly amazed how can I solve problems that I deemed unsolvable in other languages :grinning: