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,60 @@ 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
+
 ################################################################################
 
 # Adams Bashforth and Adams moulton methods
@@ -3247,7 +3307,7 @@ const MultistepAlgorithms = Union{IRKN3, IRKN4,
     ABDF2,
     AB3, AB4, AB5, ABM32, ABM43, ABM54}
 
-const SplitAlgorithms = Union{CNAB2, CNLF2, IRKC, SBDF,
+const SplitAlgorithms = Union{CNAB2, CNLF2, IRKC, SBDF, NewmarkBeta,
     KenCarp3, KenCarp4, KenCarp47, KenCarp5, KenCarp58, CFNLIRK3}
 
 #=

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))
+
+    tmp = zero(u)
+    NewmarkBetaCache(u, uprev, upred, fsalfirst, β, γ, nlsolver, tmp)
+end

src/derivative_utils.jl

@@ -629,6 +629,19 @@ function jacobian2W(mass_matrix::MT, dtgamma::Number, J::AbstractMatrix,
     return W
 end
 
+# The addition here does just work for the given example. We need to split up the jacobian into its blocks.
+function jacobian2W(mass_matrix, dtgamma::ArrayPartitionNLSolveHelper, 
+    J::AbstractMatrix, W_transform::Bool)
+    @unpack γ₁, γ₂ = dtgamma
+    jacobian2W(mass_matrix,γ₁+γ₂,J,W_transform)
+end
+
+function jacobian2W!(W::AbstractMatrix, mass_matrix, dtgamma::ArrayPartitionNLSolveHelper, 
+    J::AbstractMatrix, W_transform::Bool)
+    @unpack γ₁, γ₂ = dtgamma
+    jacobian2W!(W,mass_matrix,γ₁+γ₂,J,W_transform)
+end
+
 function calc_W!(W, integrator, nlsolver::Union{Nothing, AbstractNLSolver}, cache, dtgamma,
         repeat_step, W_transform = false, newJW = nothing)
     @unpack t, dt, uprev, u, f, p = integrator

src/integrators/integrator_utils.jl

@@ -507,6 +507,7 @@ nlsolve_f(f, alg::DAEAlgorithm) = f
 function nlsolve_f(integrator::ODEIntegrator)
     nlsolve_f(integrator.f, unwrap_alg(integrator, true))
 end
+nlsolve_f(f::DynamicalODEFunction, alg::NewmarkBeta) = f.f1 # FIXME
 
 function (integrator::ODEIntegrator)(t, ::Type{deriv} = Val{0};
         idxs = nothing) where {deriv}

src/misc_utils.jl

@@ -197,3 +197,30 @@ 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
+
+function Base.:*(γ::ArrayPartitionNLSolveHelper{T1}, z::Vector{T2}) where {T1, T2}
+    ArrayPartition(γ.γ₁*z, γ.γ₂*z)
+end
+
+Base.:*(γ::ArrayPartitionNLSolveHelper{T}, scalar::T) where T = scalar*γ
+
+function Base.:*(scalar::T, γ::ArrayPartitionNLSolveHelper{T}) where T
+    ArrayPartitionNLSolveHelper(scalar*γ.γ₁, scalar*γ.γ₂)
+end
+
+Base.inv(γ::ArrayPartitionNLSolveHelper) = 1/(γ.γ₁+γ.γ₂)

src/nlsolve/newton.jl

@@ -303,7 +303,7 @@ end
     # page 54.
     γdt = isdae ? α * invγdt : γ * dt
 
-    !(W_γdt ≈ γdt) && (rmul!(dz, 2 / (1 + γdt / W_γdt)))
+    # !(W_γdt ≈ γdt) && (rmul!(dz, 2 / (1 + γdt / W_γdt)))
     relax!(dz, nlsolver, integrator, f)
 
     calculate_residuals!(atmp, dz, uprev, ustep, opts.abstol, opts.reltol,
@@ -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)
@@ -412,6 +423,21 @@ function _compute_rhs!(tmp, ztmp, ustep, γ, α, tstep, k,
     return _vec(ztmp), ustep
 end
 
+function _compute_rhs!(tmp, ztmp, ustep, γ::ArrayPartitionNLSolveHelper, α, tstep, k,
+    invγdt, method::MethodType, p, dt, f, z)
+    mass_matrix = f.mass_matrix
+    ustep = tmp + γ * z
+    f(k, ustep, p, tstep)
+    if mass_matrix === I
+        @.. ztmp=(dt * k - z) * invγdt
+    else
+        update_coefficients!(mass_matrix, ustep, p, tstep)
+        mul!(_vec(ztmp), mass_matrix, _vec(z))
+        @.. ztmp=(dt * k - ztmp) * invγdt
+    end
+    return _vec(ztmp), ustep
+end
+
 function _compute_rhs!(tmp::Array, ztmp::Array, ustep::Array, α, tstep, k,
         invγdt, p, uprev, f::TF, z) where {TF<:DAEFunction}
     @inbounds @simd ivdep for i in eachindex(z)

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/nlsolve/type.jl

@@ -89,11 +89,11 @@ end
 
 abstract type AbstractNLSolver{algType, iip} end
 
-mutable struct NLSolver{algType, iip, uType, gamType, tmpType, tType,
+mutable struct NLSolver{algType, iip, uType, gamType, tmpType, tmp2Type, tType,
     C <: AbstractNLSolverCache} <: AbstractNLSolver{algType, iip}
     z::uType
-    tmp::uType # DIRK and multistep methods only use tmp
-    tmp2::tmpType # for GLM if neccssary
+    tmp::tmpType # DIRK and multistep methods only use tmp
+    tmp2::tmp2Type # for GLM if neccssary
     ztmp::uType
     γ::gamType
     c::tType
@@ -114,7 +114,7 @@ end
 function NLSolver{iip, tType}(z, tmp, ztmp, γ, c, α, alg, κ, fast_convergence_cutoff, ηold,
         iter, maxiters, status, cache, method = DIRK, tmp2 = nothing,
         nfails::Int = 0) where {iip, tType}
-    NLSolver{typeof(alg), iip, typeof(z), typeof(γ), typeof(tmp2), tType, typeof(cache)}(z,
+    NLSolver{typeof(alg), iip, typeof(z), typeof(γ), typeof(tmp), typeof(tmp2), tType, typeof(cache)}(z,
         tmp,
         tmp2,
         ztmp,
@@ -157,7 +157,7 @@ mutable struct NLNewtonCache{
     firststage::Bool
     firstcall::Bool
     W_γdt::tType
-    du1::uType
+    du1::rateType
     uf::ufType
     jac_config::jcType
     linsolve::lsType

src/nlsolve/utils.jl

@@ -42,6 +42,7 @@ get_new_W!(::AbstractNLSolverCache)::Bool = true
 get_W(nlsolver::AbstractNLSolver) = get_W(nlsolver.cache)
 get_W(nlcache::Union{NLNewtonCache, NLNewtonConstantCache}) = nlcache.W
 
+set_W_γdt!(nlsolver::AbstractNLSolver, W_γdt::ArrayPartitionNLSolveHelper) = set_W_γdt!(nlsolver, W_γdt.γ₁ + W_γdt.γ₂)
 set_W_γdt!(nlsolver::AbstractNLSolver, W_γdt) = set_W_γdt!(nlsolver.cache, W_γdt)
 function set_W_γdt!(nlcache::Union{NLNewtonCache, NLNewtonConstantCache}, W_γdt)
     nlcache.W_γdt = W_γdt
@@ -132,27 +133,27 @@ function build_nlsolver(alg, u, uprev, p, t, dt, f::F, rate_prototype,
         ::Type{uEltypeNoUnits},
         ::Type{uBottomEltypeNoUnits},
         ::Type{tTypeNoUnits}, γ, c,
-        iip) where {F, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+        iip; tmp = zero(u)) where {F, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
     build_nlsolver(alg, u, uprev, p, t, dt, f, rate_prototype, uEltypeNoUnits,
         uBottomEltypeNoUnits,
-        tTypeNoUnits, γ, c, 1, iip)
+        tTypeNoUnits, γ, c, 1, iip; tmp=tmp)
 end
 
 function build_nlsolver(alg, u, uprev, p, t, dt, f::F, rate_prototype,
         ::Type{uEltypeNoUnits},
         ::Type{uBottomEltypeNoUnits},
         ::Type{tTypeNoUnits}, γ, c, α,
-        iip) where {F, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+        iip; tmp = zero(u)) where {F, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
     build_nlsolver(alg, alg.nlsolve, u, uprev, p, t, dt, f, rate_prototype, uEltypeNoUnits,
-        uBottomEltypeNoUnits, tTypeNoUnits, γ, c, α, iip)
+        uBottomEltypeNoUnits, tTypeNoUnits, γ, c, α, iip; tmp=tmp)
 end
 
 function build_nlsolver(alg, nlalg::Union{NLFunctional, NLAnderson, NLNewton, NonlinearSolveAlg},
         u, uprev, p, t, dt,
         f::F, rate_prototype, ::Type{uEltypeNoUnits},
         ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits},
         γ, c, α,
-        ::Val{true}) where {F, uEltypeNoUnits, uBottomEltypeNoUnits,
+        ::Val{true}; tmp = zero(u)) where {F, uEltypeNoUnits, uBottomEltypeNoUnits,
         tTypeNoUnits}
     #TODO
     #nlalg = DiffEqBase.handle_defaults(alg, nlalg)
@@ -162,7 +163,6 @@ function build_nlsolver(alg, nlalg::Union{NLFunctional, NLAnderson, NLNewton, No
 
     # define fields of non-linear solver
     z = zero(u)
-    tmp = zero(u)
     ztmp = zero(u)
 
     # build cache of non-linear solver