{{ message }}

Instantly share code, notes, and snippets.

simonbyrne/Implicit.jl

Created Sep 1, 2021
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
 using LinearAlgebra using SparseArrays using DifferentialEquations using DiffEqOperators using BlockBandedMatrices using DiffEqOperators using RecursiveFactorization mutable struct Single_Stack L::Float64 N::Int64 t::Float64 ρ::Array{Float64,1} w::Array{Float64,1} ρθ::Array{Float64,1} ρ_res::Array{Float64,1} w_res::Array{Float64,1} ρθ_res::Array{Float64,1} _grav::Float64 γ::Float64 _R_m::Float64 Π::Function end function combine!(ρ, w, ρθ, u) N = length(ρ) u[1:N] .= ρ u[N+1:2N] .= ρθ u[2N+1:3N+1] .= w end function split!(u, ρ, w, ρθ) N = length(ρ) ρ .= u[1:N] ρθ .= u[N+1:2N] w .= u[2N+1:3N+1] end function compute_wave_speed(ρ, w, ρθ, _R_m, _C_p, γ, _MSLP) p = ( ρθ * _R_m / _MSLP^(_R_m/_C_p) ).^(γ) c = sqrt.(γ*p./ρ) return maximum(w) + maximum(c) end function DecayingTemperatureProfile(z::Float64; T_virt_surf, T_min_ref, _R_d, _grav, _MSLP) # Scale height for surface temperature H_sfc = _R_d * T_virt_surf / _grav H_t = H_sfc z′ = z / H_t tanh_z′ = tanh(z′) ΔTv = T_virt_surf - T_min_ref Tv = T_virt_surf - ΔTv * tanh_z′ ΔTv′ = ΔTv / T_virt_surf p = -H_t * (z′ + ΔTv′ * (log(1 - ΔTv′ * tanh_z′) - log(1 + tanh_z′) + z′)) p /= H_sfc * (1 - ΔTv′^2) p = _MSLP * exp(p) ρ = p/(_R_d*Tv) return (Tv, p, ρ) end function spatial_residual(du, u, ss::Single_Stack, t) @info "Tendency computation!!!!" t # u is a N+(N-1)+N vector N, L = ss.N, ss.L Δz = L/N _grav = ss._grav Π = ss.Π ρ, w, ρθ = ss.ρ, ss.w, ss.ρθ ρ_res, w_res, ρθ_res = ss.ρ_res, ss.w_res, ss.ρθ_res if eltype(u) != eltype(w) ρ = similar(ρ, eltype(u)) w = similar(w, eltype(u)) ρθ = similar(ρθ, eltype(u)) ρ_res = similar(ρ_res, eltype(u)) w_res = similar(w_res, eltype(u)) ρθ_res = similar(ρθ_res, eltype(u)) end split!(u, ρ, w, ρθ) for i = 1:N ρ₊ = (i == N ? 0.0 : (ρ[i] + ρ[i+1])/2.0) ρ₋ = (i == 1 ? 0.0 : (ρ[i] + ρ[i-1])/2.0) ρ_res[i] = -(w[i+1]ρ₊ - w[i]ρ₋)/Δz ρθ₊ = (i == N ? 0.0 : (ρθ[i] + ρθ[i+1])/2.0) ρθ₋ = (i == 1 ? 0.0 : (ρθ[i] + ρθ[i-1])/2.0) ρθ_res[i] = -(w[i+1]ρθ₊ - w[i]ρθ₋)/Δz end w_res[1] = w_res[N+1] = 0 for i = 2:N w_res[i] = -1/2*(ρθ[i-1]/ρ[i-1] + ρθ[i]/ρ[i]) *(Π(ρθ[i]) - Π(ρθ[i-1]))/Δz - _grav end combine!(ρ_res, w_res, ρθ_res, du) # @info "norm(du) = ", norm(du) end function RK4!(f, u0, ss::Single_Stack, T, Δt) u = copy(u0) k1, k2, k3, k4 = similar(u0), similar(u0), similar(u0), similar(u0) Nt = Int64(round(T / Δt)) t = 0 for i = 1:Nt f(k1, u, ss, t) f(k2, u + Δt/2*k1, ss, t+Δt/2) f(k3, u + Δt/2*k2, ss, t+Δt/2) f(k4, u + Δt*k3, ss, t+Δt) u .+= Δt/6*(k1 + 2k2 + 2k3 + k4) t += Δt end return u end function jacobian!(J,u,ss,t) @info "Jacobian computation!!!!" t # N cells N = div(length(u) - 1 , 3) ρ, ρθ, w = u[1:N], u[N+1:2N], u[2N+1:3N+1] # construct cell center ρh = [ρ[1] ; (ρ[1:N-1] + ρ[2:N])/2.0; ρ[N]] ρθh = [ρθ[1]; (ρθ[1:N-1] + ρθ[2:N])/2.0; ρθ[N]] Πc = Π.(ρθ) Πh = [NaN64; (Πc[1:N-1] + Πc[2:N])/2.0; NaN64] Δzh = [NaN64; (Δz[1:N-1] + Δz[2:N])/2.0; NaN64] for i = 1:N J[i, i+2N] = ρh[i]/Δz[i] J[i, i+2N+1] = -ρh[i+1]/Δz[i] end for i = 1:N J[i+N, i+2N] = ρθh[i]/Δz[i] J[i+N, i+2N+1] = -ρθh[i+1]/Δz[i] end # 0 for i = 1, N+1 for i = 2:N J[i+2N, (i-1)] = -_grav/(2*ρh[i]) J[i+2N, (i-1)+1] = -_grav/(2*ρh[i]) J[i+2N, (i-1)+N] = (γ - 1)*Πh[i]./(ρh[i]*Δzh[i]) J[i+2N, (i-1)+1+N] = -(γ - 1)*Πh[i]./(ρh[i]*Δzh[i]) end # D_ρ = diagm(0=>-ρh/Δz, -1=>ρh/Δz)[1:N, 1:N-1] # D_Θ = diagm(0=>-ρθh/Δz, -1=>ρθh/Δz)[1:N, 1:N-1] # G_W = (γ - 1) * diagm(0=>Πh./ρh/Δz, 1=>-Πh./ρh/Δz)[1:N-1, 1:N] # A_W = diagm(0=>-ones(N-1)./ρh/2, 1=>-ones(N-1)./ρh/2)[1:N-1, 1:N] # P = ([zeros(N,N) D_ρ zeros(N,N); # A_W*_grav zeros(N-1,N-1) G_W # zeros(N,N) D_Θ zeros(N,N)]) end function discrete_hydrostatic_balance!(ρ, w, ρθ, Δz::Float64, _grav::Float64, Π::Function) for i = 1:length(ρ)-1 ρ[i+1] = ρθ[i+1]/(-2Δz*_grav/(Π(ρθ[i+1]) - Π(ρθ[i])) - ρθ[i]/ρ[i]) end end function linsolve!(::Type{Val{:init}},f,u0;kwargs...) function _linsolve!(x,A,b,update_matrix=false;kwargs...) _A = RecursiveFactorization.lu!(A) ldiv!(x,_A,b) end end const DHB = false const ndays = 2 const N = 30 const L = 30e3 const Δz0 = L/N const _MSLP = 1e5 const _grav = 9.8 const _R_m = 287.058 const γ = 1.4 const _C_p = _R_m * γ / (γ - 1) const _C_v = _R_m / (γ - 1) const _R_d = _R_m function Π(ρθ) ρθ > 0 || return NaN _C_p*(_R_d*ρθ/_MSLP)^(_R_m/_C_v) end ρ = zeros(Float64, N) ρθ = zeros(Float64, N) w = zeros(Float64, N+1) Δz = zeros(N+1) .+ L/N for i = 1:N z = (i-0.5)*L/N Tvᵢ, pᵢ, ρᵢ = DecayingTemperatureProfile(z; T_virt_surf = 280.0, T_min_ref = 230.0, _R_d = _R_d, _grav = _grav, _MSLP = _MSLP) ρ[i] = ρᵢ ρθ[i] = ρᵢ*Tvᵢ*(_MSLP/pᵢ)^(_R_m/_C_p) end if DHB discrete_hydrostatic_balance!(ρ, w, ρθ, Δz0, _grav, Π) end ss = Single_Stack(L, N, 0.0, copy(ρ), copy(w), copy(ρθ),copy(ρ), copy(w), copy(ρθ), _grav, γ, _R_m, Π) max_wave_speed = compute_wave_speed(ρ, w, ρθ, _R_m, _C_p, γ, _MSLP) @info "maximum Δt (cfl = 1) = " Δz0 / max_wave_speed u0 = zeros(3*N+1) combine!(ρ, w, ρθ, u0) Δt = 600.0 jac_prototype = zeros(3N+1, 3N+1) jacobian!(jac_prototype,u0,ss,0.0) jac_prototype = sparse(jac_prototype) f = ODEFunction(spatial_residual, jac=jacobian!); #, jac_prototype = jac_prototype); prob_ode_hydrostatic_balance = ODEProblem(f,u0,(0.,86400 * ndays),ss); u1 = solve(prob_ode_hydrostatic_balance, Rosenbrock23(), dt = Δt, reltol=0.1) u2 = solve(prob_ode_hydrostatic_balance, Rosenbrock23(linsolve=linsolve!), dt = Δt, reltol=0.1)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters