Newmark beta method

src/OrdinaryDiffEq.jl

@@ -164,6 +164,7 @@ include("caches/prk_caches.jl")
 include("caches/pdirk_caches.jl")
 include("caches/dae_caches.jl")
 include("caches/qprk_caches.jl")
+include("caches/newmark_caches.jl")
 
 include("tableaus/low_order_rk_tableaus.jl")
 include("tableaus/high_order_rk_tableaus.jl")
@@ -213,6 +214,8 @@ include("perform_step/pdirk_perform_step.jl")
 include("perform_step/dae_perform_step.jl")
 include("perform_step/qprk_perform_step.jl")
 
+include("perform_step/newmark_perform_step.jl")
+
 include("dense/generic_dense.jl")
 include("dense/interpolants.jl")
 include("dense/rosenbrock_interpolants.jl")
@@ -404,6 +407,8 @@ export SymplecticEuler, VelocityVerlet, VerletLeapfrog, PseudoVerletLeapfrog,
 
 export SplitEuler
 
+export NewmarkBeta
+
 export Nystrom4, FineRKN4, FineRKN5, Nystrom4VelocityIndependent,
        Nystrom5VelocityIndependent,
        IRKN3, IRKN4, DPRKN4, DPRKN5, DPRKN6, DPRKN6FM, DPRKN8, DPRKN12, ERKN4, ERKN5, ERKN7, RKN4

src/alg_utils.jl

@@ -495,6 +495,7 @@ alg_order(alg::ERKN4) = 4
 alg_order(alg::ERKN5) = 5
 alg_order(alg::ERKN7) = 7
 alg_order(alg::RKN4) = 4
+alg_order(alg::NewmarkBeta) = 1 # or at least 2, depending on the parameters.
 
 alg_order(alg::Midpoint) = 2
 
@@ -1064,3 +1065,5 @@ is_mass_matrix_alg(alg::Union{OrdinaryDiffEqAlgorithm, DAEAlgorithm}) = false
 is_mass_matrix_alg(alg::CompositeAlgorithm) = all(is_mass_matrix_alg, alg.algs)
 is_mass_matrix_alg(alg::RosenbrockAlgorithm) = true
 is_mass_matrix_alg(alg::NewtonAlgorithm) = !isesdirk(alg)
+is_mass_matrix_alg(alg::NewmarkBeta) = true
+

src/algorithms.jl

@@ -47,8 +47,14 @@ abstract type DAEAlgorithm{CS, AD, FDT, ST, CJ} <: DiffEqBase.AbstractDAEAlgorit
 # Partitioned ODE Specific Algorithms
 abstract type OrdinaryDiffEqPartitionedAlgorithm <: OrdinaryDiffEqAlgorithm end
 abstract type OrdinaryDiffEqAdaptivePartitionedAlgorithm <: OrdinaryDiffEqAdaptiveAlgorithm end
+abstract type OrdinaryDiffEqImplicitPartitionedAlgorithm{CS, AD, FDT, ST, CJ} <:
+              OrdinaryDiffEqImplicitAlgorithm{CS, AD, FDT, ST, CJ} end
+abstract type OrdinaryDiffEqAdaptivePartitionedImplicitAlgorithm{CS, AD, FDT, ST, CJ} <:
+              OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS, AD, FDT, ST, CJ} end
 const PartitionedAlgorithm = Union{OrdinaryDiffEqPartitionedAlgorithm,
-    OrdinaryDiffEqAdaptivePartitionedAlgorithm}
+    OrdinaryDiffEqAdaptivePartitionedAlgorithm,
+    OrdinaryDiffEqImplicitPartitionedAlgorithm,
+    OrdinaryDiffEqAdaptivePartitionedImplicitAlgorithm}
 
 struct FunctionMap{scale_by_time} <: OrdinaryDiffEqAlgorithm end
 FunctionMap(; scale_by_time = false) = FunctionMap{scale_by_time}()
@@ -1232,6 +1238,81 @@ year = {2024},
 """
 struct RKN4 <: OrdinaryDiffEqAlgorithm end
 
+"""
+    NewmarkBeta
+
+Classical Newmark-β method to solve second order ODEs, possibly in mass matrix form.
+
+Local truncation errors are estimated with the estimate of Zienkiewicz and Xie.
+
+## References
+
+Newmark, Nathan (1959), "A method of computation for structural dynamics", 
+Journal of the Engineering Mechanics Division, 85 (EM3) (3): 67–94, doi:
+https://doi.org/10.1061/JMCEA3.0000098
+
+Zienkiewicz, O. C., and Y. M. Xie. "A simple error estimator and adaptive
+time stepping procedure for dynamic analysis." Earthquake engineering &
+structural dynamics 20.9 (1991): 871-887, doi:
+https://doi.org/10.1002/eqe.4290200907
+"""
+struct NewmarkBeta{PT, F, F2, P, CS, AD, FDT, ST, CJ} <: 
+       OrdinaryDiffEqAdaptivePartitionedImplicitAlgorithm{CS, AD, FDT, ST, CJ}
+    β::PT
+    γ::PT
+    linsolve::F
+    nlsolve::F2
+    precs::P
+end
+
+function NewmarkBeta(β, γ; chunk_size = Val{0}(), autodiff = Val{true}(), standardtag = Val{true}(),
+    concrete_jac = nothing, diff_type = Val{:forward},
+    linsolve = nothing, precs = DEFAULT_PRECS, nlsolve = NLNewton(),
+    extrapolant = :linear)
+    NewmarkBeta{
+        typeof(β), typeof(linsolve), typeof(nlsolve), typeof(precs),
+        _unwrap_val(chunk_size), _unwrap_val(autodiff), diff_type, _unwrap_val(standardtag), _unwrap_val(concrete_jac)}(
+        β, γ,
+        linsolve,
+        nlsolve,
+        precs)
+end
+
+# Needed for remake
+function NewmarkBeta(; β=nothing, γ=nothing, chunk_size = Val{0}(), autodiff = Val{true}(), standardtag = Val{true}(),
+    concrete_jac = nothing, diff_type = Val{:forward},
+    linsolve = nothing, precs = DEFAULT_PRECS, nlsolve = NLNewton(),
+    extrapolant = :linear)
+    NewmarkBeta{
+        typeof(β), typeof(linsolve), typeof(nlsolve), typeof(precs),
+        _unwrap_val(chunk_size), _unwrap_val(autodiff), diff_type, _unwrap_val(standardtag), _unwrap_val(concrete_jac)}(
+        β, γ,
+        linsolve,
+        nlsolve,
+        precs)
+end
+
+# struct CNAB2{CS, AD, F, F2, P, FDT, ST, CJ} <:
+#     OrdinaryDiffEqNewtonAlgorithm{CS, AD, FDT, ST, CJ}
+#  linsolve::F
+#  nlsolve::F2
+#  precs::P
+#  extrapolant::Symbol
+# end
+
+# function CNAB2(; chunk_size = Val{0}(), autodiff = Val{true}(), standardtag = Val{true}(),
+#      concrete_jac = nothing, diff_type = Val{:forward},
+#      linsolve = nothing, precs = DEFAULT_PRECS, nlsolve = NLNewton(),
+#      extrapolant = :linear)
+#  CNAB2{
+#      _unwrap_val(chunk_size), _unwrap_val(autodiff), typeof(linsolve), typeof(nlsolve),
+#      typeof(precs), diff_type, _unwrap_val(standardtag), _unwrap_val(concrete_jac)}(
+#      linsolve,
+#      nlsolve,
+#      precs,
+#      extrapolant)
+# end
+
 ################################################################################
 
 # Adams Bashforth and Adams moulton methods

src/caches/newmark_caches.jl

@@ -0,0 +1,32 @@
+@cache struct NewmarkBetaCache{uType, rateType, parameterType, N} <: OrdinaryDiffEqMutableCache
+    u::uType # Current solution
+    uprev::uType # Previous solution
+    upred::uType # Predictor solution
+    fsalfirst::rateType
+    β::parameterType # newmark parameter 1
+    γ::parameterType # newmark parameter 2
+    nlsolver::N # Inner solver
+    tmp::uType # temporary, because it is required.
+end
+
+function alg_cache(alg::NewmarkBeta, u, rate_prototype, ::Type{uEltypeNoUnits},
+    ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
+    dt, reltol, p, calck, 
+    ::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+
+    β = alg.β
+    γ = alg.γ
+    upred = zero(u)
+    fsalfirst = zero(rate_prototype)
+
+    @assert 0.0 ≤ β ≤ 0.5
+    @assert 0.0 ≤ γ ≤ 1.0
+
+    c = 1.0
+    γ̂ = ArrayPartitionNLSolveHelper(1.0,1.0)
+    nlsolver = build_nlsolver(alg, u.x[1], uprev.x[1], p, t, dt, f.f1, rate_prototype.x[1], uEltypeNoUnits,
+        uBottomEltypeNoUnits, tTypeNoUnits, γ̂, c, Val(true))
+
+    tmp = zero(u)
+    NewmarkBetaCache(u, uprev, upred, fsalfirst, β, γ, nlsolver, tmp)
+end

src/misc_utils.jl

@@ -197,3 +197,18 @@ function get_differential_vars(f, u)
         end
     end
 end
+
+"""
+Helper to evaluate
+    f(innertmp + γ̂⋅z, p, t + c⋅dt)
+in the nonlinear solver when z is an array partition and γ̂ is the scaling parameter.
+Note that γ̂ is different from the γ in the Newmark-β method!
+"""
+struct ArrayPartitionNLSolveHelper{T}
+    γ₁::T
+    γ₂::T
+end
+
+function Base.:*(γ::ArrayPartitionNLSolveHelper{T1}, z::ArrayPartition{T2, <: Tuple{<:Any, <:Any}}) where {T1, T2}
+    ArrayPartition(γ.γ₁*z.x[1], γ.γ₂*z.x[2])
+end

src/nlsolve/newton.jl

@@ -356,6 +356,17 @@ function compute_ustep!(ustep, tmp, γ, z, method)
     ustep
 end
 
+# Dispatch for Newmark-type solvers
+function compute_ustep!(ustep::ArrayPartition, tmp::ArrayPartition, γ::ArrayPartitionNLSolveHelper, z, method)
+    if method === COEFFICIENT_MULTISTEP
+        # ustep = z
+        @error "Multistep for second order ODEs does not work yet."
+    else
+        @.. ustep = ArrayPartition(tmp.x[1] + γ.γ₁ * z, tmp.x[1] + γ.γ₂ * z)
+    end
+    ustep
+end
+
 function _compute_rhs(tmp, γ, α, tstep, invγdt, method::MethodType, p, dt, f, z)
     mass_matrix = f.mass_matrix
     ustep = compute_ustep(tmp, γ, z, method)

src/nlsolve/nlsolve.jl

@@ -10,6 +10,9 @@ dt⋅f(innertmp + γ⋅z, p, t + c⋅dt) + outertmp = z
 ```
 
 where `dt` is the step size and `γ` and `c` are constants, and return the solution `z`.
+
+Whether `innertmp` and `outertmp` is used for the evaluation is controlled by setting `nlsolver.method`.
+In both cases the variable name is actually `nlsolver.tmp`.
 """
 function nlsolve!(nlsolver::AbstractNLSolver, integrator::DiffEqBase.DEIntegrator,
         cache = nothing, repeat_step = false)

src/perform_step/newmark_perform_step.jl

@@ -0,0 +1,78 @@
+function initialize!(integrator, cache::NewmarkBetaCache)
+    duprev, uprev = integrator.uprev.x
+    integrator.f(cache.fsalfirst, integrator.uprev, integrator.p, integrator.t)
+    integrator.fsalfirst = cache.fsalfirst
+    integrator.fsallast = du_alias_or_new(cache.nlsolver, integrator.fsalfirst)
+    integrator.stats.nf += 1
+    return
+end
+
+@muladd function perform_step!(integrator, cache::NewmarkBetaCache, repeat_step = false)
+    @unpack t, dt, f, p = integrator
+    @unpack upred, β, γ, nlsolver = cache
+
+    # Given Mu'' = f(u,u',t) we need to solve the non-linear problem
+    #   Mu(tₙ₊₁)'' - f(ũ(tₙ₊₁) + u(tₙ₊₁)'' β Δtₙ²,ũ(tₙ₊₁)' + u(tₙ₊₁)'' γ Δtₙ,t) = 0
+    # for u(tₙ₊₁)''.
+    # The predictors ũ(tₙ₊₁) are computed as
+    #   ũ(tₙ₊₁)  = u(tₙ) + u(tₙ)' Δtₙ + u(tₙ)'' (0.5 - β) Δtₙ²
+    #   ũ(tₙ₊₁)' = u(tₙ)' + u(tₙ)'' (1.0 - γ) Δtₙ
+    # such that we can compute the solution with the correctors
+    #   u(tₙ₊₁)  = ũ(tₙ₊₁)  + u(tₙ₊₁)'' β Δtₙ²
+    #   u(tₙ₊₁)' = ũ(tₙ₊₁)' + u(tₙ₊₁)'' γ Δtₙ
+
+    mass_matrix = f.mass_matrix
+
+    # Evaluate predictor
+    dduprev = integrator.fsalfirst.x[1]
+    duprev, uprev = integrator.uprev.x
+    upred_full = ArrayPartition(
+        duprev + dt*(1.0 - γ)*dduprev,
+        uprev  + dt*dt*(0.5 - β)*dduprev + dt*duprev
+    )
+
+    # _tmp = mass_matrix * @.. broadcast=false (α₁ * uprev+α₂ * uprev2)
+    # nlsolver.tmp = @.. broadcast=false _tmp/(dt * β₀)
+
+    # nlsolve!(...) solves for
+    #   dt⋅f(innertmp + γ̂⋅z, p, t + c⋅dt) + outertmp = z
+    # So we rewrite the problem
+    #     u(tₙ₊₁)'' - f(ũ(tₙ₊₁) + u(tₙ₊₁)'' β Δtₙ², ũ(tₙ₊₁)' + u(tₙ₊₁)'' γ Δtₙ,t) = 0
+    #   z = Δtₙ u(tₙ₊₁)'':
+    #     z         - Δtₙ f(ũ(tₙ₊₁) +         z β Δtₙ, ũ(tₙ₊₁)' +         z γ,t) = 0
+    #                 Δtₙ f(ũ(tₙ₊₁) +         z β Δtₙ, ũ(tₙ₊₁)' +         z γ,t) = z
+    #   γ̂ = [γ, β Δtₙ]:
+    #                 Δtₙ f(ũ(tₙ₊₁) +         z γ̂₂    , ũ(tₙ₊₁)' +         z γ̂₁   ,t) = z
+    #   innertmp = [ũ(tₙ₊₁)', ũ(tₙ₊₁)]:
+    #                 Δtₙ f(innertmp₂ +         z β Δtₙ², innertmp₁ +         z γ Δtₙ,t) = z
+    # Note: innertmp = nlsolve.tmp
+    nlsolver.γ = ArrayPartitionNLSolveHelper(γ, β * dt) # = γ̂
+    nlsolver.tmp .= upred_full # TODO check f tmp is potentially modified and if not elimiate the allocation of upred_full
+
+    # Use the linear extrapolation Δtₙ u(tₙ)'' as initial guess for the nonlinear solve
+    nlsolver.z .= dt*dduprev
+    ddu = nlsolve!(nlsolver, integrator, cache, repeat_step) / dt
+    nlsolvefail(nlsolver) && return
+
+    # Apply corrector
+    u .= ArrayPartition(
+        upred_full.x[1] + ddu*β*dt*dt,
+        upred_full.x[2] + ddu*γ*dt
+    )
+
+    #
+    if integrator.opts.adaptive
+        if integrator.success_iter == 0
+            integrator.EEst = one(integrator.EEst)
+        else
+            # Zienkiewicz and Xie (1991) Eq. 21
+            δddu = (ddu - dduprev)
+            integrator.EEst = dt*dt/2 * (2*β - 1/3) * integrator.opts.internalnorm(δddu, t)
+        end
+    end
+
+    f(integrator.fsallast, u, p, t + dt)
+    integrator.stats.nf += 1
+
+    return
+end

test/algconvergence/partitioned_methods_tests.jl

@@ -19,6 +19,7 @@ prob = DynamicalODEProblem(ff_harmonic, v0, u0, (0.0, 5.0))
 sol = solve(prob, SymplecticEuler(), dt = 1 / 2)
 sol_verlet = solve(prob, VelocityVerlet(), dt = 1 / 100)
 sol_ruth3 = solve(prob, Ruth3(), dt = 1 / 100)
+sol_newmark = solve(prob, NewmarkBeta(0.5,0.25))
 
 interp_time = 0:0.001:5
 interp = sol(0.5)
@@ -31,6 +32,7 @@ prob = SecondOrderODEProblem(f1_harmonic, v0, u0, (0.0, 5.0))
 sol2 = solve(prob, SymplecticEuler(), dt = 1 / 2)
 sol2_verlet = solve(prob, VelocityVerlet(), dt = 1 / 100)
 sol2_ruth3 = solve(prob, Ruth3(), dt = 1 / 100)
+sol2_newmark = solve(prob, NewmarkBeta(0.5,0.25))
 
 sol2_verlet(0.1)