clear code, rm ParametricFunction

src/pino_ode_solve.jl

@@ -1,8 +1,3 @@
-struct ParametricFunction{}
- function_::Union{Nothing, Function}
- bounds::Any
-end
-
 """
  PINOODE(chain,
  OptimizationOptimisers.Adam(0.1),
@@ -35,12 +30,13 @@ in which will be train the prediction of parametric ODE.
 * Sifan Wang "Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets"
 * Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations"
 """
-struct PINOODE{C, O, I, S <: Union{Nothing, AbstractTrainingStrategy},
+struct PINOODE{C, O, B, I, S <: Union{Nothing, AbstractTrainingStrategy},
  AL <: Union{Nothing, Function}, K} <:
  SciMLBase.AbstractODEAlgorithm
  chain::C
  opt::O
- parametric_function::ParametricFunction
+ bounds::B
+ number_of_parameters::Int
  init_params::I
  strategy::S
  additional_loss::AL
@@ -49,13 +45,14 @@ end
 
 function PINOODE(chain,
  opt,
- parametric_function;
+ bounds,
+ number_of_parameters;
  init_params = nothing,
  strategy = nothing,
  additional_loss = nothing,
  kwargs...)
  !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain))
- PINOODE(chain, opt, parametric_function, init_params, strategy, additional_loss, kwargs)
+ PINOODE(chain, opt, bounds,number_of_parameters, init_params, strategy, additional_loss, kwargs)
 end
 
 mutable struct PINOPhi{C, S}
@@ -82,64 +79,34 @@ function (f::PINOPhi{C, T})(x::NamedTuple, θ) where {C <: NeuralOperator, T}
  y
 end
 
-function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C <: DeepONet, T}
- # @unpack function_, bounds = parametric_function
- # branch_left, branch_right = function_.(p, t), function_.(p, t .+ sqrt(eps(eltype(t))))
- pfs, p, t = x
+function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T}
+ p, t = x
  branch_left, branch_right = p, p
  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 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 physics_loss(
  phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T}
- pfs, p, t = x
+ p, t = x
  f = prob.f
  tuple = (branch = p, trunk = t)
  out = phi(tuple, θ)
  if size(p)[1] == 1
  fs = f.(out, p, t)
  f_vec= vec(fs)
- # out_ = vec(out)
  else
  f_vec = reduce(vcat,[[f(out[i], p[:, i, 1], t[j]) for i in axes(p, 2)] for j in axes(t, 3)])
  end
- du = vec(dfdx(phi, x, θ, prob))
+ du = vec(dfdx(phi, x, θ))
  norm = prod(size(du))
  sum(abs2, du .- f_vec) / norm
 end
 
 function initial_condition_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T}
- pfs, p, t = x
+ p, t = x
  t0 = t[:, :, [1]]
  # pfs0 = pfs[:, :, [1]]
  tuple = (branch = p, trunk = t0)
@@ -150,40 +117,26 @@ function initial_condition_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) whe
  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, parametric_function::ParametricFunction, tspan)
- @unpack function_, bounds = parametric_function
+function get_trainset(strategy::GridTraining, bounds, number_of_parameters, tspan)
  dt = strategy.dx
- #TODO
- size_of_p = 50
- if bounds isa Tuple
- p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2])
+ if size(bounds)[1] == 1
+ bound = bounds[1]
+ p_ = range(start = bound[1], length = number_of_parameters, stop = bound[2])
  p = collect(reshape(p_, 1, size(p_)[1], 1))
  else
- p_ = [range(start = b[1], length = size_of_p, stop = b[2]) for b in bounds]
+ p_ = [range(start = b[1], length = number_of_parameters, stop = b[2])
+ for b in bounds]
  p = vcat([collect(reshape(p_i, 1, size(p_i)[1], 1)) for p_i in p_]...)
  end
 
  t_ = collect(tspan[1]:dt:tspan[2])
  t = reshape(t_, 1, 1, size(t_)[1])
- pfs = function_.(p,t)
- (pfs, p, t)
+ (p, t)
 end
 
 function generate_loss(
- strategy::GridTraining, prob::ODEProblem, phi, parametric_function::ParametricFunction, tspan)
- x = get_trainset(strategy, parametric_function, tspan)
+ strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan)
+ x = get_trainset(strategy, bounds, number_of_parameters, tspan)
  function loss(θ, _)
  initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ)
  end
@@ -198,7 +151,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem,
  saveat = nothing,
  maxiters = nothing)
  @unpack tspan, u0, f = prob
- @unpack chain, opt, parametric_function, init_params, strategy, additional_loss = alg
+ @unpack chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss=alg
 
  if !isa(chain, DeepONet)
  error("Only DeepONet neural networks are supported")
@@ -231,7 +184,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem,
  throw(ArgumentError("Only GridTraining strategy is supported"))
  end
 
- inner_f = generate_loss(strategy, prob, phi, parametric_function, tspan)
+ inner_f = generate_loss(strategy, prob, phi, bounds, number_of_parameters, tspan)
 
  function total_loss(θ, _)
  L2_loss = inner_f(θ, nothing)
@@ -241,10 +194,6 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem,
  L2_loss
  end
 
- # TODO delete
- # total_loss(θ, 0)
- # Zygote.gradient(θ -> total_loss(θ, 0), θ)
-
  # Optimization Algo for Training Strategies
  opt_algo = Optimization.AutoZygote()
 
@@ -261,7 +210,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem,
  optprob = OptimizationProblem(optf, init_params)
  res = solve(optprob, opt; callback, maxiters, alg.kwargs...)
 
- pfs, p, t = get_trainset(strategy, parametric_function, tspan)
+ p, t = get_trainset(strategy, bounds, number_of_parameters, tspan)
  tuple = (branch = p, trunk = t)
  u = phi(tuple, res.u)
 

test/PINO_ode_tests.jl

@@ -1,5 +1,5 @@
 using Test
-using OrdinaryDiffEq, OptimizationOptimisers
+using OptimizationOptimisers
 using Lux
 using Statistics, Random
 using NeuralPDE
@@ -27,24 +27,22 @@ using NeuralPDE
  θ, st = Lux.setup(Random.default_rng(), deeponet)
 
  c = deeponet(x, θ, st)[1]
- function_(p, t) = cos(p*t)
- bounds = (pi, 2pi)
- parametric_function = ParametricFunction(function_, bounds)
+ bounds = [(pi, 2pi)]
+ number_of_parameters = 50
  dt = (tspan[2] - tspan[1]) / 40
  strategy = GridTraining(dt)
  opt = OptimizationOptimisers.Adam(0.01)
- alg = PINOODE(deeponet, opt, parametric_function; strategy = strategy)
+ alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy)
  sol = solve(prob, alg, verbose = true, maxiters = 3000)
 
- phi(tuple, sol.original.u)
  sol.original.objective
  # TODO intrepretation output another mesh
  # x = (branch = p, trunk = t)
  # phi(sol.original.u)
  # sol.
  ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p)
- size_of_p = 50
- p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2])
+ #TDOD another number_of_parameters
+ p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2])
  p = collect(reshape(p_, 1, size(p_)[1], 1))
  ground_solution = ground_analytic.(u0, p, sol.t.trunk)
 
@@ -66,23 +64,21 @@ end
  Lux.Dense(10, 10, Lux.tanh_fast))
 
  deeponet = DeepONet(branch, trunk)
- function_(p, t) = cos(p * t)
- bounds = (0.1f0, 2.f0)
- parametric_function = ParametricFunction(function_, bounds)
- sol.original.objective
+ bounds = [(0.1f0, 2.f0)]
+ number_of_parameters = 40
  dt = (tspan[2] - tspan[1]) / 40
  strategy = GridTraining(dt)
 
  opt = OptimizationOptimisers.Adam(0.01)
- alg = PINOODE(deeponet, opt, parametric_function; strategy = strategy)
+ alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy)
 
- sol = solve(prob, alg, verbose = false, maxiters = 5000)
+ sol = solve(prob, alg, verbose = false, maxiters = 3000)
  sol.original.objective
  #if u0 == 1
  ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) /
  (p^2 + 1)
- size_of_p = 50
- p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2])
+
+ p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2])
  p = collect(reshape(p_, 1, size(p_)[1], 1))
  ground_solution = ground_analytic_.(u0, p, sol.t.trunk)
 
@@ -106,11 +102,8 @@ end
  linear = Lux.Chain(Lux.Dense(10, 1))
  deeponet = DeepONet(branch, trunk; linear = linear)
 
- function_(p, t) = cos(p * t)
- bounds = (0.0f0, 10.0f0)
- parametric_function = ParametricFunction(function_, bounds)
-
- # db = (bounds.p[2] - bounds.p[1]) / 50
+ bounds = [(0.0f0, 10.0f0)]
+ number_of_parameters = 60
  dt = (tspan[2] - tspan[1]) / 40
  strategy = GridTraining(dt)
 
@@ -133,24 +126,23 @@ end
  end
 
  branch_size, trunk_size = 50, 40
- p,t = get_trainset(branch_size, trunk_size, bounds, tspan)
- data, tuple = get_data()
+ p,t = get_trainset(branch_size, trunk_size, bounds[1], tspan)
+ data, tuple_ = get_data()
 
  function additional_loss_(phi, θ)
- u = phi(tuple, θ)
+ u = phi(tuple_, θ)
  norm = prod(size(u))
  sum(abs2, u .- data) / norm
  end
  alg = PINOODE(
- deeponet, opt, parametric_function; strategy = strategy, additional_loss = additional_loss_)
+ deeponet, opt, bounds, number_of_parameters; strategy = strategy,
+ additional_loss = additional_loss_)
  sol = solve(prob, alg, verbose = true, maxiters = 2000)
-
- size_of_p = 50
- p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2])
+ p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][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.005
+ @test ground_solution≈sol.u rtol=0.01
 end
 
 #vector outputs and multiple parameters
@@ -175,38 +167,21 @@ end
  Lux.Dense(10, 10, Lux.tanh_fast))
 
  deeponet = DeepONet(branch, trunk)
-
- #TODO add size_of_p = 50
- function_(p, t) = cos(p * t)
  bounds = [(0.1f0, pi), (1.0f0, 2.0f0)]
- parametric_function = ParametricFunction(function_, bounds)
+ number_of_parameters = 50
  dt = (tspan[2] - tspan[1]) / 40
  strategy = GridTraining(dt)
  opt = OptimizationOptimisers.Adam(0.03)
- alg = PINOODE(deeponet, opt, parametric_function; strategy = strategy)
+ alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy)
  sol = solve(prob, alg, verbose = true, maxiters = 3000)
 
  ga = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t)
- p_ = [range(start = b[1], length = size_of_p, stop = b[2]) for b in bounds]
+ p_ = [range(start = b[1], length = number_of_parameters, 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 = sol.t.trunk
  ground_solution_ = f_vec = reduce(
  hcat, [reduce(
  vcat, [ga(u0, p[:, i, 1], t[j]) for i in axes(p, 2)]) for j in axes(t, 3)])
  ground_solution = reshape(ground_solution_, 1, size(ground_solution_)...)
  @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 = 10)
-# end
-
-# plot_()