Skip to content

Commit

Permalink
energies: adding energy loss functions
Browse files Browse the repository at this point in the history 
This commit adds energy loss functions. These are loss functions
analogous to those obtained from PDE equations or boundary equations,
but the energy integrand can be given explicitly in symbolic form.
  • Loading branch information
@drsk0
drsk0 committed Sep 12, 2023
1 parent af4d4b7 commit 90b8a40
Show file tree
Hide file tree
Showing 7 changed files with 420 additions and 138 deletions.
64 changes: 46 additions & 18 deletions src/adaptive_losses.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ end
"""
```julia
NonAdaptiveLoss{T}(; pde_loss_weights = 1,
energy_loss_weights = 1,
bc_loss_weights = 1,
additional_loss_weights = 1)
```
Expand All @@ -25,30 +26,33 @@ change during optimization
mutable struct NonAdaptiveLoss{T <: Real} <: AbstractAdaptiveLoss
pde_loss_weights::Vector{T}
bc_loss_weights::Vector{T}
energy_loss_weights::Vector{T}
additional_loss_weights::Vector{T}
SciMLBase.@add_kwonly function NonAdaptiveLoss{T}(; pde_loss_weights = 1,
energy_loss_weights = 1,
bc_loss_weights = 1,
additional_loss_weights = 1) where {
T <:
Real
}
new(vectorify(pde_loss_weights, T), vectorify(bc_loss_weights, T),
new(vectorify(pde_loss_weights, T), vectorify(energy_loss_weights, T), vectorify(bc_loss_weights, T),
vectorify(additional_loss_weights, T))
end
end

# default to Float64
SciMLBase.@add_kwonly function NonAdaptiveLoss(; pde_loss_weights = 1, bc_loss_weights = 1,
SciMLBase.@add_kwonly function NonAdaptiveLoss(; pde_loss_weights = 1, energy_loss_weights = 1, bc_loss_weights = 1,
additional_loss_weights = 1)
NonAdaptiveLoss{Float64}(; pde_loss_weights = pde_loss_weights,
bc_loss_weights = bc_loss_weights,
energy_loss_weights = energy_loss_weights,
additional_loss_weights = additional_loss_weights)
end

function generate_adaptive_loss_function(pinnrep::PINNRepresentation,
adaloss::NonAdaptiveLoss,
pde_loss_functions, bc_loss_functions)
function null_nonadaptive_loss(θ, pde_losses, bc_losses)
pde_loss_functions, energy_loss_functions, bc_loss_functions)
function null_nonadaptive_loss(θ, pde_loss, energy_loss, bc_losses)
nothing
end
end
Expand All @@ -58,6 +62,7 @@ end
GradientScaleAdaptiveLoss(reweight_every;
weight_change_inertia = 0.9,
pde_loss_weights = 1,
energy_loss_weights = 1,
bc_loss_weights = 1,
additional_loss_weights = 1)
```
Expand Down Expand Up @@ -90,61 +95,66 @@ mutable struct GradientScaleAdaptiveLoss{T <: Real} <: AbstractAdaptiveLoss
reweight_every::Int64
weight_change_inertia::T
pde_loss_weights::Vector{T}
energy_loss_weights::Vector{T}
bc_loss_weights::Vector{T}
additional_loss_weights::Vector{T}
SciMLBase.@add_kwonly function GradientScaleAdaptiveLoss{T}(reweight_every;
weight_change_inertia = 0.9,
pde_loss_weights = 1,
energy_loss_weights = 1,
bc_loss_weights = 1,
additional_loss_weights = 1) where {
T <:
Real
}
new(convert(Int64, reweight_every), convert(T, weight_change_inertia),
vectorify(pde_loss_weights, T), vectorify(bc_loss_weights, T),
vectorify(additional_loss_weights, T))
vectorify(pde_loss_weights, T), vectorify(energy_loss_weights, T),
vectorify(bc_loss_weights, T), vectorify(additional_loss_weights, T))
end
end
# default to Float64
SciMLBase.@add_kwonly function GradientScaleAdaptiveLoss(reweight_every;
weight_change_inertia = 0.9,
pde_loss_weights = 1,
energy_loss_weights = 1,
bc_loss_weights = 1,
additional_loss_weights = 1)
GradientScaleAdaptiveLoss{Float64}(reweight_every;
weight_change_inertia = weight_change_inertia,
pde_loss_weights = pde_loss_weights,
energy_loss_weights = energy_loss_weights,
bc_loss_weights = bc_loss_weights,
additional_loss_weights = additional_loss_weights)
end

function generate_adaptive_loss_function(pinnrep::PINNRepresentation,
adaloss::GradientScaleAdaptiveLoss,
pde_loss_functions, bc_loss_functions)
pde_loss_functions, energy_loss_functions, bc_loss_functions)
weight_change_inertia = adaloss.weight_change_inertia
iteration = pinnrep.iteration
adaloss_T = eltype(adaloss.pde_loss_weights)

function run_loss_gradients_adaptive_loss(θ, pde_losses, bc_losses)
if iteration[1] % adaloss.reweight_every == 0
# the paper assumes a single pde loss function, so here we grab the maximum of the maximums of each pde loss function
pde_grads_maxes = [maximum(abs.(Zygote.gradient(pde_loss_function, θ)[1]))
for pde_loss_function in pde_loss_functions]
pde_grads_max = maximum(pde_grads_maxes)
# we treat energy loss functions the same as pde loss functions
pde_energy_grads_maxes = [maximum(abs.(Zygote.gradient(pde_loss_function, θ)[1]))
for pde_loss_function in vcat(pde_loss_functions, energy_loss_functions)]
pde_energy_grads_max = maximum(pde_energy_grads_maxes)
bc_grads_mean = [mean(abs.(Zygote.gradient(bc_loss_function, θ)[1]))
for bc_loss_function in bc_loss_functions]

nonzero_divisor_eps = adaloss_T isa Float64 ? Float64(1e-11) :
convert(adaloss_T, 1e-7)
bc_loss_weights_proposed = pde_grads_max ./
bc_loss_weights_proposed = pde_energy_grads_max ./
(bc_grads_mean .+ nonzero_divisor_eps)
adaloss.bc_loss_weights .= weight_change_inertia .*
adaloss.bc_loss_weights .+
(1 .- weight_change_inertia) .*
bc_loss_weights_proposed
logscalar(pinnrep.logger, pde_grads_max, "adaptive_loss/pde_grad_max",
logscalar(pinnrep.logger, pde_energy_grads_max, "adaptive_loss/pde_energy_grad_max",
iteration[1])
logvector(pinnrep.logger, pde_grads_maxes, "adaptive_loss/pde_grad_maxes",
logvector(pinnrep.logger, pde_energy_grads_maxes, "adaptive_loss/pde_energy_grad_maxes",
iteration[1])
logvector(pinnrep.logger, bc_grads_mean, "adaptive_loss/bc_grad_mean",
iteration[1])
Expand All @@ -160,8 +170,10 @@ end
```julia
function MiniMaxAdaptiveLoss(reweight_every;
pde_max_optimiser = Flux.ADAM(1e-4),
energy_max_optimiser = Flux.ADAM(1e-4),
bc_max_optimiser = Flux.ADAM(0.5),
pde_loss_weights = 1,
energy_loss_weights = 1,
bc_loss_weights = 1,
additional_loss_weights = 1)
```
Expand Down Expand Up @@ -191,65 +203,81 @@ https://arxiv.org/abs/2009.04544
"""
mutable struct MiniMaxAdaptiveLoss{T <: Real,
PDE_OPT <: Flux.Optimise.AbstractOptimiser,
ENERGY_OPT <: Flux.Optimise.AbstractOptimiser,
BC_OPT <: Flux.Optimise.AbstractOptimiser} <:
AbstractAdaptiveLoss
reweight_every::Int64
pde_max_optimiser::PDE_OPT
energy_max_optimiser::ENERGY_OPT
bc_max_optimiser::BC_OPT
pde_loss_weights::Vector{T}
energy_loss_weights::Vector{T}
bc_loss_weights::Vector{T}
additional_loss_weights::Vector{T}
SciMLBase.@add_kwonly function MiniMaxAdaptiveLoss{T,
PDE_OPT, BC_OPT}(reweight_every;
PDE_OPT, ENERGY_OPT, BC_OPT}(reweight_every;
pde_max_optimiser = Flux.ADAM(1e-4),
energy_max_optimiser = Flux.ADAM(1e-4),
bc_max_optimiser = Flux.ADAM(0.5),
pde_loss_weights = 1,
energy_loss_weights = 1,
bc_loss_weights = 1,
additional_loss_weights = 1) where {
T <:
Real,
PDE_OPT <:
Flux.Optimise.AbstractOptimiser,
ENERGY_OPT <:
Flux.Optimise.AbstractOptimiser,
BC_OPT <:
Flux.Optimise.AbstractOptimiser
}
new(convert(Int64, reweight_every), convert(PDE_OPT, pde_max_optimiser),
new(convert(Int64, reweight_every), convert(PDE_OPT, pde_max_optimiser), convert(ENERGY_OPT, energy_max_optimiser),
convert(BC_OPT, bc_max_optimiser),
vectorify(pde_loss_weights, T), vectorify(bc_loss_weights, T),
vectorify(additional_loss_weights, T))
vectorify(energy_loss_weights, T), vectorify(additional_loss_weights, T))
end
end

# default to Float64, ADAM, ADAM
SciMLBase.@add_kwonly function MiniMaxAdaptiveLoss(reweight_every;
pde_max_optimiser = Flux.ADAM(1e-4),
energy_max_optimiser = Flux.ADAM(1e-4),
bc_max_optimiser = Flux.ADAM(0.5),
pde_loss_weights = 1,
energy_loss_weights = 1,
bc_loss_weights = 1,
additional_loss_weights = 1)
MiniMaxAdaptiveLoss{Float64, typeof(pde_max_optimiser),
typeof(bc_max_optimiser)}(reweight_every;
pde_max_optimiser = pde_max_optimiser,
energy_max_optimiser = energy_max_optimiser,
bc_max_optimiser = bc_max_optimiser,
pde_loss_weights = pde_loss_weights,
energy_loss_weights = energy_loss_weights,
bc_loss_weights = bc_loss_weights,
additional_loss_weights = additional_loss_weights)
end

function generate_adaptive_loss_function(pinnrep::PINNRepresentation,
adaloss::MiniMaxAdaptiveLoss,
pde_loss_functions, bc_loss_functions)
pde_loss_functions, energy_loss_functions, bc_loss_functions)
pde_max_optimiser = adaloss.pde_max_optimiser
energy_max_optimiser = adaloss.energy_max_optimiser
bc_max_optimiser = adaloss.bc_max_optimiser
iteration = pinnrep.iteration

function run_minimax_adaptive_loss(θ, pde_losses, bc_losses)
function run_minimax_adaptive_loss(θ, pde_losses, energy_losses, bc_losses)
if iteration[1] % adaloss.reweight_every == 0
Flux.Optimise.update!(pde_max_optimiser, adaloss.pde_loss_weights,
-pde_losses)
Flux.Optimise.update!(energy_max_optimiser, adaloss.energy_loss_weights,
-energy_losses)
Flux.Optimise.update!(bc_max_optimiser, adaloss.bc_loss_weights, -bc_losses)
logvector(pinnrep.logger, adaloss.pde_loss_weights,
"adaptive_loss/pde_loss_weights", iteration[1])
logvector(pinnrep.logger, adaloss.energy_loss_weights,
"adaptive_loss/energy_loss_weights", iteration[1])
logvector(pinnrep.logger, adaloss.bc_loss_weights,
"adaptive_loss/bc_loss_weights",
iteration[1])
Expand Down
Loading

0 comments on commit 90b8a40

Please sign in to comment.