some minor tweaks towards making n-body example working again

Author: lecopivo

SciLean/Analysis/Calculus/HasRevFDeriv.lean

@@ -1013,6 +1013,36 @@ theorem Norm2.norm2.arg_a0.HasRevFDerivUpdate_simple_rule :
   case adjoint => intro; dsimp; data_synth
   case simp => funext x; simp; funext dr x'; module
 
+set_option linter.unusedVariables false in
+@[data_synth]
+theorem norm2.arg_a0.HasRevFDeriv_rule
+    (f : X → Y) (f')
+    (hf : HasRevFDeriv R f f') (hf' : ∀ x, f x ≠ 0) :
+    HasRevFDeriv R (fun x => ‖f x‖₂[R])
+      (fun x =>
+        let' (y,df) := f' x
+        let ynorm := ‖y‖₂[R]
+        (ynorm, fun dr =>
+          let dy := (dr * ynorm⁻¹) • y
+          let dx := df dy
+          dx)) := by
+  sorry_proof
+
+set_option linter.unusedVariables false in
+@[data_synth]
+theorem norm2.arg_a0.HasRevFDerivUpdate_rule
+    (f : X → Y) (f')
+    (hf : HasRevFDerivUpdate R f f') (hf' : ∀ x, f x ≠ 0) :
+    HasRevFDerivUpdate R (fun x => ‖f x‖₂[R])
+      (fun x =>
+        let' (y,df) := f' x
+        let ynorm := ‖y‖₂[R]
+        (ynorm, fun dr dx =>
+          let dy := (dr * ynorm⁻¹) • y
+          let dx := df dy dx
+          dx)) := by
+  sorry_proof
+
 end OverReals
 
 

SciLean/Analysis/ODE/OdeSolve.lean

@@ -26,6 +26,16 @@ structure HasUniqueOdeSolution (f : R → X → X) extends HasOdeSolution f : Pr
   uniq : ∀ t₀ x₀ x x', IsOdeSolution f t₀ x₀ x → IsOdeSolution f t₀ x₀ x' → x = x'
 
 open Classical in
+/-- Solution of ordinary differentiale equation.
+
+Function `x := fun t => odeSolve f t₀ t x₀` satisfies ODE
+```
+∂ x t = f t (x t)
+```
+with initial condition `x t₀ = x₀`.
+
+Junk value is returned if `f` does define ODE with an unique solution.
+ -/
 noncomputable
 def odeSolve (f : R → X → X) (t₀ t : R) (x₀ : X) : X :=
   if h : HasUniqueOdeSolution f -- TODO: can we reduce it to just HasOdeSolution?

SciLean/Numerics/Optimization/Optimjl.lean

@@ -41,37 +41,3 @@ theorem argmin_eq_limit_optimize
       let f' := holdLet <| revFDeriv R f
       let r := optimize {f:=f,f':=f',hf:=sorry_proof} (AbstractOptimizer.setOptions X method opts) x₀
       r.minimizer := sorry_proof
-
-
-#check LBFGS Float 1
-
-set_default_scalar Float
-
-open Scalar
-
-approx mySqrt_v1 (y : Float) := argmin (x : Float), ‖x^2-y‖₂²
-by
-  rw[argmin_eq_limit_optimize (method := (default : LBFGS Float 1)) (x₀:=1)]
-
-  approx_limit options sorry_proof
-
-  conv in (revFDeriv _ _) => autodiff
-
-
-def mySqrt_v2  (y : Float) :=
-  let r := optimize (d:={f := fun (x : Float) => ‖x^2-y‖₂²,
-                         f' := _,
-                         hf := by data_synth => enter[3]; lsimp})
-                    (method := (default : LBFGS Float 1))
-                    (x₀ := 1)
-  r.minimizer
-
-
-
-/-- info: 1.414212 -/
-#guard_msgs in
-#eval mySqrt_v1 2 ({},())
-
-/-- info: 1.414214 -/
-#guard_msgs in
-#eval mySqrt_v1 2 ({x_abstol := 1e-16, g_abstol := 0},())

SciLean/Tactic/DataSynth/Elab.lean

@@ -33,7 +33,7 @@ syntax (name:=data_synth_conv) "data_synth" optConfig (discharger)? : conv
     | none => fun _ => return none
     | some stx => Mathlib.Meta.FunProp.tacticToDischarge ⟨stx.raw[3]⟩
 
-  let (r?,_) ← dataSynth g |>.run {config := cfg, discharge := disch} |>.run stateRef
+  let (r?,_) ← dataSynth g |>.run {config := cfg, discharge := fun e => do disch e} |>.run stateRef
     |>.run (← Simp.mkDefaultMethods).toMethodsRef
     |>.run (← Simp.mkContext (config := cfg.toConfig) (simpTheorems := #[← getSimpTheorems]))
     |>.run {}
@@ -81,7 +81,7 @@ syntax (name:=data_synth_tac) "data_synth" optConfig (discharger)? ("=>" convSeq
 
   let stateRef : IO.Ref DataSynth.State ← IO.mkRef {}
 
-  let (r?,_) ← dataSynth g |>.run {config := cfg, discharge := disch} |>.run stateRef
+  let (r?,_) ← dataSynth g |>.run {config := cfg, discharge := fun e => do disch e} |>.run stateRef
     |>.run (← Simp.mkDefaultMethods).toMethodsRef
     |>.run (← Simp.mkContext (config := cfg.toConfig) (simpTheorems := #[← getSimpTheorems]))
     |>.run {}

SciLean/Tactic/DataSynth/Types.lean

@@ -120,7 +120,7 @@ structure Config extends DataSynthConfig, Simp.Config
 structure Context where
   config : Config := {}
   normalize : Expr → Simp.SimpM Simp.Result := fun e => return {expr := e}
-  discharge : Expr → MetaM (Option Expr) := fun _ => return .none
+  discharge : Expr → SimpM (Option Expr) := fun _ => return .none
 
 structure State where
   numSteps := 0

Test/data_synth/get_elem.lean

@@ -0,0 +1,41 @@
+import SciLean.Data.ArrayOperations.Operations.GetElem
+import SciLean.Data.DataArray.Float
+import SciLean.Data.DataArray.VectorType
+
+open SciLean
+
+variable {n} (i : Fin n) (j : Fin 3)
+
+/--
+info: HasRevFDeriv Float (fun x => x[i]) fun x =>
+  (x[i], fun xi =>
+    let x' := setElem 0 i xi True.intro;
+    x') : Prop
+-/
+#guard_msgs in
+#check (HasRevFDeriv Float (fun x : Float^[n] => x[i]) _ ) rewrite_by data_synth
+
+
+/--
+info: HasRevFDeriv Float (fun x => x[(i, j)]) fun x =>
+  (x[(i, j)], fun xi =>
+    let x' := setElem 0 (i, j) xi True.intro;
+    x') : Prop
+-/
+#guard_msgs in
+#check (HasRevFDeriv Float (fun x : Float^[n,3] => x[i,j]) _ ) rewrite_by data_synth
+
+
+/--
+info: HasRevFDeriv Float (fun x => x[(i, j)]) fun x =>
+  (x[(i, j)], fun xi =>
+    let x' := setElem 0 (i, j) xi True.intro;
+    x') : Prop
+-/
+#guard_msgs in
+#check (HasRevFDeriv Float (fun x : Float^[3]^[n] => x[i,j]) _ ) rewrite_by data_synth
+
+
+-- this is broken!!!
+-- some serious issue with type classes :(
+-- #check (HasRevFDeriv Float (fun x : Float^[3]^[n] => x[i]) _ ) rewrite_by data_synth

examples/NBody.lean

@@ -1,39 +1,180 @@
-import SciLean.Mechanics
-import SciLean.Operators.ODE
-import SciLean.Solver 
-import SciLean.Tactic.LiftLimit
-import SciLean.Tactic.FinishImpl
-import SciLean.Data.PowType
-import SciLean.Core.Extra
-import SciLean.Functions
+import SciLean
+
+import ProofWidgets.Data.Svg
+import ProofWidgets.Component.HtmlDisplay
+
+open Lean ProofWidgets
+
 
 open SciLean
 
 set_option synthInstance.maxSize 2048
 set_option synthInstance.maxHeartbeats 500000
 set_option maxHeartbeats 500000
 
-def H (n : Nat) (ε : ℝ) (m : Idx n → ℝ) (x p : ((ℝ^(3:Nat))^n)) : ℝ := 
-  (∑ i, (1/(2*m i)) * ∥p[i]∥²)
-  +
-  (∑ i j, (m i*m j) * ∥x[i] - x[j]∥^{(-1:ℝ),ε})
-argument p [Fact (ε≠0)] [Fact (n≠0)]
-  isSmooth, hasAdjDiff, adjDiff
-argument x [Fact (ε≠0)] [Fact (n≠0)]
-  isSmooth, hasAdjDiff,
-  adjDiff by
-    simp[H]
-    simp [sum_into_lambda]
-    simp [← sum_of_add]
-    
-def solver (steps : ℕ) (n : Nat) [Fact (n≠0)] (ε : ℝ) [Fact (ε≠0)] (m : Idx n → ℝ)
-  : Impl (ode_solve (HamiltonianSystem (H n ε m))) :=
+macro (priority:=high) A:term noWs "[" i:term "," ":" "]" : term => `(MatrixType.row $A $i)
+macro (priority:=high) A:term noWs "[" ":" "," j:term "]" : term => `(MatrixType.col $A $j)
+
+axiom unsafeNonzero {α} [Zero α] (a : α) : a ≠ 0
+
+macro "unsafeAD" : tactic =>
+  `(tactic| (intros; simp only [not_false_eq_true, ne_eq, unsafeNonzero]))
+
+
+open Lean Meta
+instance : MonadLift Tactic.DataSynth.DataSynthM SimpM where
+  monadLift e := do
+    let disch? := (← Simp.getMethods).discharge?
+    -- discharge? : Expr → SimpM (Option Expr) := fun _ => return none
+    let r := e { discharge := disch? } (← ST.mkRef {}) (← ST.mkRef {})
+    r
+
+
+theorem revFDeriv_from_hasRevFDeriv {K} [RCLike K]
+  {X} [NormedAddCommGroup X] [AdjointSpace K X]
+  {Y} [NormedAddCommGroup Y] [AdjointSpace K Y]
+  {f : X → Y} {f'} (hf : HasRevFDeriv K f f') :
+  revFDeriv K f = f' := sorry_proof
+
+
+open Lean Meta in
+/-- Compute `revFDeriv R f` with calling data_synth on `HasRevFDeriv R f ?f'`. -/
+simproc_decl revFDeriv_simproc (revFDeriv _ _) := fun e => do
+
+  -- get field and function to differentiate
+  let K := e.getArg! 0
+  let f := e.appArg!
+
+  -- craft `HasRevFDeriv K f ?f'`
+  let goal ← mkAppM ``HasRevFDeriv #[K,f]
+  let (xs,_,_) ← forallMetaTelescope (← inferType goal)
+  let f' := xs[0]!
+  let goal := goal.app f'
+
+  -- extract data_synth goal
+  let .some goal ← Tactic.DataSynth.isDataSynthGoal? goal
+    | throwError "something went really wrong"
+
+  -- run data_synth
+  let .some r ← Tactic.DataSynth.dataSynth goal
+    | return .continue
+
+  let f'' := r.xs[0]!
+  let prf ← mkAppM ``revFDeriv_from_hasRevFDeriv #[r.proof]
+  -- IO.println (← ppExpr (← inferType prf))
+
+  -- let eq ← mkEq e f''
+  -- let prf ← mkSorry eq false
+
+  return .visit { expr := f'', proof? := prf }
+
+
+set_default_scalar Float
+open Scalar
+example : (<∂ x : Float, x) = fun x => (x, fun dx => dx) := by simp[revFDeriv_simproc]
+example : (<∂ x : Float, x*x) = fun x => (x*x, fun dx => x*dx+x*dx) := by simp[revFDeriv_simproc]
+example : (<∂ x : Float, exp x) = fun x => (exp x, fun dx => dx*exp x) := by simp[revFDeriv_simproc]
+
+def H {n} (m : Fin n → Float) (x p : Float^[n,2]) : Float :=
+  (∑ i, (1/(2*m i)) * ‖p[i,:]‖₂²)
+  -
+  (∑ (i : Fin n) (j : Fin i.1),
+    let j := j.castLT (by omega)
+    (m i*m j) * ‖x[i,:] - x[j,:]‖₂⁻¹)
+
+
+
+#check (<∂ x':=(⊞[1.0,2.0,3.0]:Float3), ‖x'‖₂⁻¹) rewrite_by
+  unfold fgradient
+  lsimp (disch:=unsafeAD) only [simp_core,revFDeriv_simproc]
+
+variable (x p : Float^[n,3]) (ε : Float)
+
+-- set_option trace.Meta.Tactic.data_synth true in
+-- #check (<∂ x':=x, ∑ i j, ‖x'[i,:] - x'[j,:]‖₂⁻¹) rewrite_by
+--   unfold fgradient
+--   lsimp (disch:=unsafeAD) only [simp_core, revFDeriv_simproc]
+
+-- #check (<∂ x':=x, H (fun _ => 1) x' p) rewrite_by
+--   unfold fgradient H
+--   lsimp (disch:=unsafeAD) only [simp_core, revFDeriv_simproc]
+
+-- #check (<∂ p':=p, H (fun _ => 1) x p') rewrite_by
+--   unfold fgradient H
+--   lsimp (disch:=unsafeAD) only [simp_core, revFDeriv_simproc]
+
+
+-- #check odeSolve.arg_x₀.revFDeriv_rule
+
+
+approx solver (m : Fin n → Float)
+  := odeSolve (fun (t : Float) (x,p) => ( ∇ (p':=p), H m x  p',
+                                         -∇ (x':=x), H m x' p))
 by
-  -- Unfold Hamiltonian definition and compute gradients
-  simp[HamiltonianSystem]
-  
+  -- Unfold Hamiltonian and compute gradients
+  unfold H fgradient
+  lsimp (disch:=unsafeAD) only [simp_core,revFDeriv_simproc]
+
   -- Apply RK4 method
-  rw [ode_solve_fixed_dt runge_kutta4_step]
-  lift_limit steps "Number of ODE solver steps."; admit; simp
-   
-  finish_impl
+  simp_rw (config:={zeta:=false}) [odeSolve_fixed_dt rungeKutta4 sorry_proof]
+
+  -- todo: make approx_limit ignore leading let bindings
+  approx_limit n sorry_proof
+
+
+
+#eval! solver (fun _ : Fin 2 => 1) (10,()) 0 0.7
+         (⊞[-1.0,0;1.0,0],⊞[0.0,-1.0;0,1.0])
+
+
+def generateData (m : ℕ) : Float^[2,2]^[m] := Id.run do
+
+  let mut data : Float^[2,2]^[m] := 0
+  let Δt := 0.01
+
+  -- initial state
+  let mut x := ⊞[-1.0, 0.0;
+                  1.0, 0.0]
+
+  let mut p := ⊞[ 0.0,-0.3;
+                  0.0, 0.3]
+
+  for h : i in [0:m] do
+    let i : Fin m := ⟨i, sorry_proof⟩
+    data[i] := x
+    (x,p) := solver 1 (1,()) 0 Δt (x,p)
+
+  return data
+
+
+open ProofWidgets Svg
+
+private def frame : Frame where
+  xmin   := -2
+  ymin   := -2
+  xSize  := 4
+  width  := 400
+  height := 400
+
+open IndexType
+instance {I} [IndexType I] : ToJson (Float^[I]) where
+  toJson x := toJson (Array.ofFn (fun i => x[fromFin i]))
+
+instance {I} [IndexType I] : FromJson (Float^[I]) where
+  fromJson? j :=
+    match fromJson? (α:=Array Float) j with
+    | Except.error e => Except.error e
+    | Except.ok data =>
+      Except.ok (⊞ (i : I) => data[toFin i]!)
+
+
+instance (f : Frame) : Coe (Float^[2]) (Point f) := ⟨fun x => .abs x[0] x[1]⟩
+
+private def svg : Svg frame :=
+  let data := generateData 1000
+
+  { elements :=
+      #[polyline (Array.ofFn (fun i => data[i][0,:])) |>.setStroke (0.8,0.8,0.2) (.px 1) ,
+        polyline (Array.ofFn (fun i => data[i][1,:])) |>.setStroke (0.2,0.8,0.8) (.px 1) ] }
+
+#html svg.toHtml

examples/WaveEquation.lean

@@ -1,5 +1,4 @@
 import SciLean
-import SciLean
 
 open SciLean