Merge pull request #242 from YichengDWu/addmacro

Add @showprogress macro

Project.toml

@@ -24,6 +24,8 @@ NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+ProgressBars = "49802e3a-d2f1-5c88-81d8-b72133a6f568"
 QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
@@ -63,6 +65,7 @@ NNlibCUDA = "0.2"
 Optimisers = "0.2"
 Optimization = "3"
 OptimizationOptimisers = "0.1"
+ProgressBars = "1.5" 
 QuasiMonteCarlo = "0.2, 0.3"
 Requires = "1"
 RuntimeGeneratedFunctions = "0.5"

docs/make.jl

@@ -15,7 +15,7 @@ makedocs(; plugins=[bib], modules=[Sophon], sitename="Sophon.jl",
                                 canonical="https://YichengDWu.github.io/Sophon.jl",
                                 edit_link="main", assets=String[indigo]),
          authors="Yicheng Wu",
-         warnonly = Documenter.except(:autodocs_block, :cross_references, :docs_block, 
+         warnonly = Documenter.except(:autodocs_block, :cross_references, :docs_block,
                                       :doctest, :eval_block, :example_block, :footnote,
                                       :linkcheck_remotes, :linkcheck, :meta_block, :parse_error),
          pages=[

docs/src/qa.md

@@ -16,8 +16,9 @@ residual = Sophon.residual_function_1(prob.p[1], res.u)
 If you want to monitor the loss during training, create a callback function like the following:
 ```julia
 function callback(p, _)
-    loss = sum(abs2, Sophon.residual_function_1(prob.p[1], p))
-    println("loss: $loss")
+    loss_1 = sum(abs2, Sophon.residual_function_1(prob.p[1], p))
+    loss_2 = sum(abs2, Sophon.residual_function_1(prob.p[2], p))
+    println("loss: $loss_1, loss_2")
     return false
 end
 ```

docs/src/tutorials/L_shape.md

@@ -38,7 +38,7 @@ strategy = NonAdaptiveTraining()
 
 prob = Sophon.discretize(pde_system, pinn, sampler, strategy)
 
-res = Optimization.solve(prob, BFGS(); maxiters=1000)
+@showprogress res = Optimization.solve(prob, BFGS(); maxiters=1000)
 
 using CairoMakie
 

docs/src/tutorials/SchrödingerEquation.md

@@ -52,7 +52,7 @@ function train(pde_system, prob, sampler, strategy, resample_period = 500, n=10)
      return res
 end
 
-res = train(pde_system, prob, sampler, strategy)
+@showprogress res = train(pde_system, prob, sampler, strategy)
 ```
 
 ```@example Schrödinger

docs/src/tutorials/allen_cahn.md

@@ -6,7 +6,7 @@ In this tutorial we are going to solve the Allen-Cahn equation with periodic bou
 ```@example allen
 using ModelingToolkit, IntervalSets
 using Sophon
-using Optimization, OptimizationOptimJL
+using Optimization, OptimizationOptimJL, Zygote
 
 @parameters t, x
 @variables u(..)

docs/src/tutorials/convection.md

@@ -47,7 +47,7 @@ pinn = PINN(chain)
 
 prob = Sophon.discretize(convection, pinn, sampler, strategy) 
 
-@time res = Optimization.solve(prob, BFGS(); maxiters = 1000)
+@showprogress res = Optimization.solve(prob, BFGS(); maxiters = 1000)
 ```
 
 Let's visualize the result.

docs/src/tutorials/helmholtz.md

@@ -44,7 +44,7 @@ strategy = NonAdaptiveTraining()
 
 prob = Sophon.discretize(helmholtz, pinn, sampler, strategy) 
 
-@time res = Optimization.solve(prob, BFGS(); maxiters=1000)
+@showprogress res = Optimization.solve(prob, BFGS(); maxiters=1000)
 ```
 
 Let's plot the result.
@@ -71,3 +71,5 @@ fig
 save("helmholtz.png", fig); nothing # hide
 ```
 ![](helmholtz.png)
+
+

docs/src/tutorials/inverse.md

@@ -69,11 +69,7 @@ function additional_loss(phi, θ)
 end
 prob = Sophon.discretize(pde_system, pinn, sampler, strategy, additional_loss=additional_loss)
 
-callback = function (p,l)
-    println("Current loss is: $l")
-    return false
-end
-@time res = Optimization.solve(prob, BFGS(), callback = callback, maxiters=1000)
+@showprogress res = Optimization.solve(prob, BFGS(), maxiters=1000)
 ```
 
 ```julia

docs/src/tutorials/ode.md

@@ -42,7 +42,7 @@ strategy = NonAdaptiveTraining(1,50)
 prob = Sophon.discretize(lotka_volterra, pinn, sampler, strategy)
 
 # Solving the problem using BFGS optimization
-res = Optimization.solve(prob, BFGS(); maxiters=1000)
+@showprogress res = Optimization.solve(prob, BFGS(); maxiters=1000)
 ```
 Next, we'll compare our results with a reference solution to verify our computations.
 

docs/src/tutorials/poisson.md

@@ -30,7 +30,7 @@ sampler = QuasiRandomSampler(200, 1)
 strategy = NonAdaptiveTraining(1 , 50)
 
 prob = Sophon.discretize(poisson, pinn, sampler, strategy)
-res = Optimization.solve(prob, BFGS(); maxiters=2000)
+@showprogress res = Optimization.solve(prob, BFGS(); maxiters=2000)
 
 phi = pinn.phi
 xs = 0:0.001:1

docs/src/tutorials/sod.md

@@ -56,7 +56,7 @@ end
 strategy = AdaptiveTraining(pde_weights, Returns(10))
 prob = Sophon.discretize(pde_system, pinn, sampler, strategy)
 
-res = Optimization.solve(prob, BFGS(); maxiters=1000)
+@showprogress res = Optimization.solve(prob, BFGS(); maxiters=1000)
 
 θ = res.u
 phi = pinn.phi

docs/src/tutorials/waveinverse2.md

@@ -56,7 +56,7 @@ additional_loss(phi, θ) = sum(abs2, phi.u(input_data, θ.u) .- u_data)
 
 prob = Sophon.discretize(wave, pinn, sampler, strategy; additional_loss=additional_loss)
 
-@time res = Optimization.solve(prob, BFGS(), maxiters=1000)
+@showprogress res = Optimization.solve(prob, BFGS(), maxiters=1000)
 ```
 
 Let's visualize the predictted solution and inferred velocity

src/Sophon.jl

@@ -30,6 +30,7 @@ using StaticArrays: SVector
 
 RuntimeGeneratedFunctions.init(@__MODULE__)
 
+include("showprogress.jl")
 include("layers/basic.jl")
 include("layers/activations.jl")
 include("layers/containers.jl")
@@ -52,6 +53,7 @@ function __init__()
     end
 end
 
+export @showprogress
 export gaussian, quadratic, laplacian, expsin, multiquadratic, stan
 export FourierFeature, TriplewiseFusion, FullyConnected, ResNet, Sine, RBF,
        DiscreteFourierFeature, ConstantFunction, ScalarLayer, SplitFunction, FactorizedDense

src/showprogress.jl

@@ -0,0 +1,52 @@
+using MacroTools
+using ProgressBars
+using Printf
+
+macro showprogress(expr)
+    @capture(expr, res_ = Optimization.solve(prob_, method_; kwargs__)) ||
+        error("Unsupported expression format")
+
+    maxiters = nothing
+    callback = nothing
+
+    for kw in kwargs
+        isnothing(maxiters) && @capture(kw, maxiters=maxiters_)
+        isnothing(callback) && @capture(kw, callback=callback_)
+    end
+
+    if maxiters === nothing
+        error("maxiters is not specified")
+    end
+
+    iter_var = gensym(:iter)
+    progress_expr = :($iter_var = ProgressBars.ProgressBar(1:$maxiters))
+
+    if isnothing(callback)
+        callback_var = gensym(:callback)
+        callback_def = quote
+            function $callback_var(p, l)
+                ProgressBars.update($iter_var)
+                ProgressBars.set_description($iter_var, string(Printf.@sprintf("%.6e", l)))
+                return false
+            end
+        end
+    else
+        callback_var = gensym(:callback)
+        callback_def = quote
+            function $callback_var(p, l)
+                update($iter_var)
+                set_description($iter_var, string(@sprintf("%.6e", l)))
+                return $callback(p, l)
+            end
+        end
+    end
+
+    new_func_call = Expr(:call, :(Optimization.solve), esc(prob), esc(method), kwargs..., :($(Expr(:kw, :callback, callback_var))))
+    assignment_expr = Expr(:(=), esc(res), new_func_call)
+
+    return quote
+        $progress_expr
+        $callback_def
+        $assignment_expr
+    end
+end