Skip to content

Commit

Permalink
Out-of-place implementation of PDE solver (Issue ODINN-SciML#21)
Browse files Browse the repository at this point in the history
  • Loading branch information
facusapienza21 committed Aug 8, 2023
1 parent 9c6d28c commit 63dd5ad
Show file tree
Hide file tree
Showing 6 changed files with 228 additions and 70 deletions.
1 change: 1 addition & 0 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ authors = ["Jordi Bolibar <[email protected]>"]
version = "0.1.0"

[deps]
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
Infiltrator = "5903a43b-9cc3-4c30-8d17-598619ec4e9b"
JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
Expand Down
34 changes: 33 additions & 1 deletion src/models/iceflow/SIA2D/SIA2D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ include("SIA2D_utils.jl")
mutable struct SIA2Dmodel{F <: AbstractFloat, I <: Integer} <: SIAmodel
A::Union{Ref{F}, Nothing}
n::Union{Ref{F}, Nothing}
H₀::Union{Matrix{F}, Nothing}
H::Union{Matrix{F}, Nothing}
::Union{Matrix{F}, Nothing}
S::Union{Matrix{F}, Nothing}
Expand Down Expand Up @@ -40,6 +41,7 @@ end
function SIA2Dmodel(params::Sleipnir.Parameters;
A::Union{Ref{F}, Nothing} = nothing,
n::Union{Ref{F}, Nothing} = nothing,
H₀::Union{Matrix{F}, Nothing} = nothing,
H::Union{Matrix{F}, Nothing} = nothing,
::Union{Matrix{F}, Nothing} = nothing,
S::Union{Matrix{F}, Nothing} = nothing,
Expand Down Expand Up @@ -68,7 +70,7 @@ function SIA2Dmodel(params::Sleipnir.Parameters;

ft = params.simulation.float_type
it = params.simulation.int_type
SIA2D_model = SIA2Dmodel{ft,it}(A, n, H, H̄, S, dSdx, dSdy, D, Dx, Dy, dSdx_edges, dSdy_edges,
SIA2D_model = SIA2Dmodel{ft,it}(A, n, H₀, H, H̄, S, dSdx, dSdy, D, Dx, Dy, dSdx_edges, dSdy_edges,
∇S, ∇Sx, ∇Sy, Fx, Fy, Fxx, Fyy, V, Vx, Vy, Γ, MB, MB_mask, MB_total, glacier_idx)

return SIA2D_model
Expand Down Expand Up @@ -99,6 +101,7 @@ function initialize_iceflow_model!(iceflow_model::IF,
F = params.simulation.float_type
iceflow_model.A = Ref{F}(glacier.A)
iceflow_model.n = Ref{F}(glacier.n)
iceflow_model.H₀ = deepcopy(glacier.H₀)::Matrix{F}
iceflow_model.H = deepcopy(glacier.H₀)::Matrix{F}
iceflow_model.= zeros(F,nx-1,ny-1)
iceflow_model.S = deepcopy(glacier.S)::Matrix{F}
Expand Down Expand Up @@ -126,5 +129,34 @@ function initialize_iceflow_model!(iceflow_model::IF,
iceflow_model.glacier_idx = Ref{I}(glacier_idx)
end

"""
function initialize_iceflow_model(iceflow_model::IF,
glacier_idx::I,
glacier::AbstractGlacier,
params::Sleipnir.Parameters
) where {IF <: IceflowModel, I <: Int}
Initialize iceflow model data structures to enable out-of-place mutation.
Keyword arguments
=================
- `iceflow_model`: Iceflow model used for simulation.
- `glacier_idx`: Index of glacier.
- `glacier`: `Glacier` to provide basic initial state of the ice flow model.
- `parameters`: `Parameters` to configure some physical variables.
"""
function initialize_iceflow_model(iceflow_model::IF,
glacier_idx::I,
glacier::Sleipnir.AbstractGlacier,
params::Sleipnir.Parameters
) where {IF <: IceflowModel, I <: Int}
nx, ny = glacier.nx, glacier.ny
F = params.simulation.float_type
iceflow_model.A = Ref{F}(glacier.A)
iceflow_model.n = Ref{F}(glacier.n)
iceflow_model.glacier_idx = Ref{I}(glacier_idx)
# We just need initial condition to run out-of-place forward model
iceflow_model.H₀ = deepcopy(glacier.H₀)::Matrix{F}
iceflow_model.H = deepcopy(glacier.H₀)::Matrix{F}
end

92 changes: 49 additions & 43 deletions src/models/iceflow/SIA2D/SIA2D_utils.jl
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@


"""
SIA!(dH, H, SIA2Dmodel)
SIA2D!(dH, H, SIA2Dmodel)
Compute an in-place step of the Shallow Ice Approximation PDE in a forward model
"""
Expand Down Expand Up @@ -84,29 +84,37 @@ function noSIA2D!(dH::Matrix{F}, H::Matrix{F}, simulation::SIM, t::F) where {F <
end

"""
SIA(H, A, context)
SIA(H, SIA2Dmodel)
Compute a step of the Shallow Ice Approximation UDE in a forward model. Allocates memory.
"""
function SIA2D(H, SIA2Dmodel)
function SIA2D(H::Matrix{F}, simulation::SIM, t::F) where {F <: AbstractFloat, SIM <: Simulation}
# function SIA2D(H, SIA2Dmodel)
# Retrieve parameters
B = SIA2Dmodel.B
Δx = SIA2Dmodel.Δx
Δy = SIA2Dmodel.Δy
A = SIA2Dmodel.A
n = SIA2Dmodel.n
SIA2D_model::SIA2Dmodel = simulation.model.iceflow
glacier::Sleipnir.Glacier2D = simulation.glaciers[simulation.model.iceflow.glacier_idx[]]
params::Sleipnir.Parameters = simulation.parameters
int_type = simulation.parameters.simulation.int_type
# Retrieve parameters
B::Matrix{F} = glacier.B
Δx::F = glacier.Δx
Δy::F = glacier.Δy
A::Ref{F} = SIA2D_model.A
n::Ref{F} = SIA2D_model.n
ρ::F = params.physical.ρ
g::F = params.physical.g

# Update glacier surface altimetry
S = B .+ H

# All grid variables computed in a staggered grid
# Compute surface gradients on edges
dSdx = diff_x(S) ./ Δx
dSdy= diff_y(S) ./ Δy
∇S = (avg_y(dSdx).^2 .+ avg_x(dSdy).^2).^((n - 1)/2)
dSdy = diff_y(S) ./ Δy
∇S = (avg_y(dSdx).^2 .+ avg_x(dSdy).^2).^((n[] - 1)/2)

Γ = 2.0 * A ** g)^n / (n+2) # 1 / m^3 s
D = Γ .* avg(H).^(n + 2) .* ∇S
Γ = 2.0 * A[] ** g)^n[] / (n[]+2) # 1 / m^3 s
D = Γ .* avg(H).^(n[] + 2) .* ∇S

# Compute flux components
@views dSdx_edges = diff_x(S[:,2:end - 1]) ./ Δx
Expand All @@ -120,10 +128,6 @@ function SIA2D(H, SIA2Dmodel)
dSdx_edges .= @views @. max(dSdx_edges, -η₀ * H[1:end-1, 2:end-1]/Δx)
dSdy_edges .= @views @. min(dSdy_edges, η₀ * H[2:end-1, 2:end]/Δy)
dSdy_edges .= @views @. max(dSdy_edges, -η₀ * H[2:end-1, 1:end-1]/Δy)
# dSdx_edges .= min.(dSdx_edges, η₀ * H[1:end-1, 2:end-1]./Δx, η₀ * H[2:end, 2:end-1]./Δx)
# dSdy_edges .= min.(dSdy_edges, η₀ * H[2:end-1, 1:end-1]./Δy, η₀ * H[2:end-1, 2:end]./Δy)
# dSdx_edges .= max.(dSdx_edges, -η₀ * H[1:end-1, 2:end-1]./Δx, -η₀ * H[2:end, 2:end-1]./Δx)
# dSdy_edges .= max.(dSdy_edges, -η₀ * H[2:end-1, 1:end-1]./Δy, -η₀ * H[2:end-1, 2:end]./Δy)

Fx = .-avg_y(D) .* dSdx_edges
Fy = .-avg_x(D) .* dSdy_edges
Expand All @@ -135,21 +139,22 @@ function SIA2D(H, SIA2Dmodel)
end

"""
avg_surface_V(context, H, temp, sim, θ=[], UA_f=[])
avg_surface_V(simulation::SIM)
Computes the average ice surface velocity for a given glacier evolution period
based on the initial and final ice thickness states.
"""
function avg_surface_V!(H₀::Matrix{F}, simulation::SIM) where {F <: AbstractFloat, SIM <: Simulation}
# We compute the initial and final surface velocity and average them
function avg_surface_V!(simulation::SIM) where {F <: AbstractFloat, SIM <: Simulation}
# TODO: Add more datapoints to better interpolate this
ft = simulation.parameters.simulation.float_type
iceflow_model = simulation.model.iceflow
Vx₀::Matrix{ft}, Vy₀::Matrix{ft} = surface_V!(H₀, simulation)
Vx::Matrix{ft}, Vy::Matrix{ft} = surface_V!(iceflow_model.H, simulation)

Vx₀::Matrix{ft}, Vy₀::Matrix{ft} = surface_V!(iceflow_model.H₀, simulation)
Vx::Matrix{ft}, Vy::Matrix{ft} = surface_V!(iceflow_model.H, simulation)

iceflow_model.Vx .= (Vx₀ .+ Vx)./2.0
iceflow_model.Vy .= (Vy₀ .+ Vy)./2.0
iceflow_model.V .= (iceflow_model.Vx.^2 .+ iceflow_model.Vy.^2).^(1/2)
iceflow_model.V .= (iceflow_model.Vx.^2 .+ iceflow_model.Vy.^2).^(1/2)
end

"""
Expand All @@ -158,19 +163,20 @@ end
Computes the average ice surface velocity for a given glacier evolution period
based on the initial and final ice thickness states.
"""
function avg_surface_V(ifm::IF, temp::F, sim, θ=nothing, UA_f=nothing; testmode=false) where {F <: AbstractFloat, IF <: IceflowModel}
# TODO: see how to get this
# A_noise = parameters.A_noise
# Δx, Δy, A_noise = retrieve_context(context, sim)
function avg_surface_V(H::Matrix{F}, simulation::SIM) where {F <: AbstractFloat, SIM <: Simulation}

ft = simulation.parameters.simulation.float_type
iceflow_model = simulation.model.iceflow

# We compute the initial and final surface velocity and average them
# TODO: Add more datapoints to better interpolate this
Vx₀, Vy₀ = surface_V(ifm, sim, A_noise, θ, UA_f; testmode=testmode)
Vx, Vy = surface_V(ifm, temp, sim, A_noise, θ, UA_f; testmode=testmode)
Vx₀::Matrix{ft}, Vy₀::Matrix{ft} = surface_V(iceflow_model.H₀, simulation)
Vx::Matrix{ft}, Vy::Matrix{ft} = surface_V(H, simulation)

V̄x = (Vx₀ .+ Vx)./2.0
V̄y = (Vy₀ .+ Vy)./2.0
V = (V̄x.^2 .+ V̄y.^2).^(1/2)

return V̄x, V̄y
return V̄x, V̄y, V

end

Expand Down Expand Up @@ -230,12 +236,18 @@ end
Computes the ice surface velocity for a given glacier state
"""
function surface_V(SIA2Dmodel, temp, sim, A_noise, θ=nothing, UA_f=nothing; testmode=false)

B = SIA2Dmodel.B
H = SIA2Dmodel.H
Δx = SIA2Dmodel.Δx
Δy = SIA2Dmodel.Δy
function surface_V(H::Matrix{F}, simulation::SIM) where {F <: AbstractFloat, SIM <: Simulation}
params::Sleipnir.Parameters = simulation.parameters
ft = params.simulation.float_type
iceflow_model = simulation.model.iceflow
glacier::Sleipnir.Glacier2D = simulation.glaciers[iceflow_model.glacier_idx[]]
B = glacier.B
Δx = glacier.Δx
Δy = glacier.Δy
A::Ref{ft} = iceflow_model.A
n::Ref{ft} = iceflow_model.n
ρ::ft = params.physical.ρ
g::ft = params.physical.g

# Update glacier surface altimetry
S = B .+ H
Expand All @@ -246,13 +258,7 @@ function surface_V(SIA2Dmodel, temp, sim, A_noise, θ=nothing, UA_f=nothing; tes
dSdy = diff_y(S) / Δy
∇S = (avg_y(dSdx).^2 .+ avg_x(dSdy).^2).^((n[] - 1)/2)

@assert (sim == "UDE" || sim == "PDE" || sim == "UDE_inplace") "Wrong type of simulation. Needs to be 'UDE' , 'UDE_inplace' or 'PDE'."
if (sim == "UDE" && !testmode) || (sim == "UDE_inplace" && !testmode)
A = predict_A̅(UA_f, θ, [temp])[1]
elseif sim == "PDE" || testmode
A = A_fake(temp, A_noise, noise)[1]
end
Γꜛ = 2.0 * A * (ρ[] * g[])^n[] / (n[]+1) # surface stress (not average) # 1 / m^3 s
Γꜛ = 2.0 * A[] ** g)^n[] / (n[]+1) # surface stress (not average) # 1 / m^3 s
D = Γꜛ .* avg(H).^(n[] + 1) .* ∇S

# Compute averaged surface velocities
Expand Down
Loading

0 comments on commit 63dd5ad

Please sign in to comment.