update multiple parameters task

src/pino_ode_solve.jl

@@ -80,18 +80,26 @@ end
 function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C <: DeepONet, T}
  p, t = x
  f = prob.f
- # if any(in(keys(bounds)), (:u0,))
- # branch_left, branch_right = f.(0, zeros(size(p)), t .+ sqrt(eps(eltype(p)))),
- # f.(0, zeros(size(p)), t)
- # else
- branch_left, branch_right = f.(0, p, t .+ sqrt(eps(eltype(p)))), f.(0, p, t)
- # end
+ branch_left, branch_right = f.(0, p, t .+ sqrt(eps(eltype(p)))), f.(0, p, t)
  trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t
  x_left = (branch = branch_left, trunk = trunk_left)
  x_right = (branch = branch_right, trunk = trunk_right)
  (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(t)))
 end
 
+# function physics_loss(
+# phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T}
+# p, t = x
+# f = prob.f
+# du = vec(dfdx(phi, x, θ, prob))
+# f_ = f.(0, p, t)
+# tuple = (branch = f_, trunk = t)
+# out = phi(tuple, θ)
+# f_ = vec(f.(out, p, t))
+# norm = prod(size(out))
+# sum(abs2, du .- f_) / norm
+# end
+
 function physics_loss(
  phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T}
  p, t = x
@@ -101,50 +109,69 @@ function physics_loss(
  # tuple = (branch = p, trunk = t)
  # out = phi(tuple, θ)
  # f.(out, p, t) ?
+ #TODO if DeepONet else err
  du = vec(dfdx(phi, x, θ, prob))
- # if any(in(keys(bounds)), (:u0,))
- # f_ = f.(0, zeros(size(p)), t)
- # else
- f_ = f.(0, p, t)
- # end
- #TODO if DeepONet else error
- tuple = (branch = f_, trunk = t)
+ fs_ = hcat([hcat([f(0, p[:, i, :], t[j]) for i in axes(p, 2)]) for j in axes(t, 3)]...)
+ fs = reshape(fs_, (1, size(fs_)...))
+ tuple = (branch = p, trunk = t)
  out = phi(tuple, θ)
- f_ = vec(f.(out, p, t))
+
+ tuple = (branch = fs, trunk = t) #TODO -> tuple = (branch = p, trunk = t)
+ out = phi(tuple, θ)
+ # f_ = vec(f.(out, p, t))
  norm = prod(size(out))
  sum(abs2, du .- f_) / norm
 end
 
-function inital_condition_loss(
+# function initial_condition_loss(
+# phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T}
+# p, t = x
+# f = prob.f
+# t0 = t[:, :, [1]]
+# f_0 = f.(0, p, t0)
+# tuple = (branch = f_0, trunk = t0)
+# out = phi(tuple, θ)
+# u = vec(out)
+# u0_ = fill(prob.u0, size(out))
+# u0 = vec(u0_)
+
+# norm = prod(size(u0))
+# sum(abs2, u .- u0) / norm
+# end
+
+function initial_condition_loss(
  phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T}
  p, t = x
  f = prob.f
- t0 = t[:, :, [1]]
- #TODO
- # if any(in(keys(bounds)), (:u0,))
- # u0_ = collect(p)
- # u0 = vec(u0_)
- # #TODO f_0 = f.(0, p, t0) and p call as u0
- # f_0 = f.(0, zeros(size(u0_)), t0)
- # else
- u0_ = fill(prob.u0, size(out))
- u0 = vec(u0_)
- f_0 = f.(0, p, t0)
- # end
- tuple = (branch = f_0, trunk = t0)
+ t0 = t[:, :, [5]]
+ fs_0 = hcat([f(0, p[:, i, :], t[j]) for i in axes(p, 2)]...)
+ fs0 = reshape(fs_0, (1, size(fs_0)...))
+ tuple = (branch = fs0, trunk = t0)
  out = phi(tuple, θ)
  u = vec(out)
+ u0_ = fill(prob.u0, size(out))
+ u0 = vec(u0_)
 
  norm = prod(size(u0))
  sum(abs2, u .- u0) / norm
 end
 
+# function get_trainset(strategy::GridTraining, bounds, tspan)
+# db, dt = strategy.dx
+# v = values(bounds)[1]
+# #TODO for all v
+# p_ = v[1]:db:v[2]
+# p = reshape(p_, 1, size(p_)[1], 1)
+# t_ = collect(tspan[1]:dt:tspan[2])
+# t = reshape(t_, 1, 1, size(t_)[1])
+# (p, t)
+# end
+
 function get_trainset(strategy::GridTraining, bounds, tspan)
- db, dt = strategy.dx
- v = values(bounds)[1]
- #TODO for all v
- p_ = v[1]:db:v[2]
- p = reshape(p_, 1, size(p_)[1], 1)
+ dt = strategy.dx
+ size_of_p = 50
+ p_ = [range(start = b[1], length = size_of_p, stop = b[2]) for b in bounds]
+ p = vcat([collect(reshape(p_i, 1, size(p_i)[1], 1)) for p_i in p_]...)
  t_ = collect(tspan[1]:dt:tspan[2])
  t = reshape(t_, 1, 1, size(t_)[1])
  (p, t)
@@ -174,9 +201,10 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem,
  # if !any(in(keys(bounds)), (:p, :u0))
  # error("bounds should contain p only")
  # end
- if !any(in(keys(bounds)), (:p,))
- error("bounds should contain p only")
- end
+ #TODO new p
+ # if !any(in(keys(bounds)), (:p,))
+ # error("bounds should contain p only")
+ # end
  phi, init_params = generate_pino_phi_θ(chain, init_params)
 
  isinplace(prob) &&
@@ -194,7 +222,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem,
  end
 
  if strategy isa GridTraining
- if length(strategy.dx) !== 2
+ if length(strategy.dx) !== 2 #TODO ?
  throw(ArgumentError("The discretization should have two elements dx= [db,dt],
  steps for branch and trunk bounds"))
  end

test/PINO_ode_tests.jl

@@ -28,6 +28,7 @@ using NeuralPDE
  c = deeponet(x, θ, st)[1]
 
  bounds = (p = [0.1f0, pi],)
+ # bounds = [0.1f0, pi]
  db = (bounds.p[2] - bounds.p[1]) / 50
  dt = (tspan[2] - tspan[1]) / 40
  strategy = GridTraining([db, dt])
@@ -141,25 +142,14 @@ end
  @test ground_solution≈sol.u rtol=0.005
 end
 
-plot(sol.u[1, :, :], linetype = :contourf)
-plot!(ground_solution[1, :, :], linetype = :contourf)
-
-function plot_()
- # Animate
- anim = @animate for (i) in 1:41
- plot(ground_solution[1, i, :], label = "Ground")
- # plot(equation_[1, i, :], label = "equation")
- plot!(sol.u[1, i, :], label = "Predicted")
- end
- gif(anim, "pino.gif", fps = 15)
-end
-
-plot_()
-
 #vector outputs and multiple parameters
 @testset "Example du = cos(p * t)" begin
- # equation = (u, p, t) -> cos(p1 * t) + p2
- equation = (u, p, t) -> cos(p[1] * t) + p[2]
+ function equation1(u, p, t)
+ p1, p2 = p[1], p[2]
+ cos(p1 * t) + p2
+ end
+
+ equation = (u, p, t) -> cos(p * t)
  tspan = (0.0f0, 1.0f0)
  u0 = 1.0f0
  prob = ODEProblem(equation, u0, tspan)
@@ -173,20 +163,39 @@ plot_()
  Lux.Dense(10, 10, Lux.tanh_fast),
  Lux.Dense(10, 10, Lux.tanh_fast))
 
- deeponet = NeuralPDE.DeepONet(branch, trunk; linear = nothing)
-
- bounds = (p1 = [0.1f0, pi], p2 = [0.1f0, 2.0f0])
- db = (bounds.u0[2] - bounds.u0[1]) / 50
+ deeponet = DeepONet(branch, trunk; linear = nothing)
+ # p1 = [0.1f0, pi]; p2 = [0.1f0, 2.0f0]
+ # bounds = (p = [p1, p2],)
+ #TODO add size_of_p = 50
+ bounds = [[0.1f0, pi], [0.1f0, 2.0f0]]
+ # db = 0.025f0
  dt = (tspan[2] - tspan[1]) / 40
- strategy = NeuralPDE.GridTraining([db, dt])
+ strategy = GridTraining(dt)
  opt = OptimizationOptimisers.Adam(0.03)
- alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy)
+ alg = PINOODE(deeponet, opt, bounds; strategy = strategy)
  sol = solve(prob, alg, verbose = false, maxiters = 2000)
- ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p)
 
+ ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p)
  p_ = bounds.p[1]:strategy.dx[1]:bounds.p[2]
  p = reshape(p_, 1, size(p_)[1], 1)
  ground_solution = ground_analytic.(u0, p, sol.t.trunk)
 
  @test ground_solution≈sol.u rtol=0.01
 end
+
+
+
+plot(sol.u[1, :, :], linetype = :contourf)
+plot!(ground_solution[1, :, :], linetype = :contourf)
+
+function plot_()
+ # Animate
+ anim = @animate for (i) in 1:41
+ plot(ground_solution[1, i, :], label = "Ground")
+ # plot(equation_[1, i, :], label = "equation")
+ plot!(sol.u[1, i, :], label = "Predicted")
+ end
+ gif(anim, "pino.gif", fps = 15)
+end
+
+plot_()