Skip to content

Commit

Permalink
test: first order tests
Browse files Browse the repository at this point in the history
  • Loading branch information
avik-pal committed Oct 30, 2024
1 parent 52a2387 commit d0389b7
Show file tree
Hide file tree
Showing 26 changed files with 855 additions and 695 deletions.
71 changes: 71 additions & 0 deletions .github/workflows/CI_NonlinearSolveFirstOrder.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
name: CI (NonlinearSolveFirstOrder)

on:
pull_request:
branches:
- master
paths:
- "lib/NonlinearSolveFirstOrder/**"
- ".github/workflows/CI_NonlinearSolveFirstOrder.yml"
- "lib/NonlinearSolveBase/**"
- "lib/SciMLJacobianOperators/**"
push:
branches:
- master

concurrency:
# Skip intermediate builds: always.
# Cancel intermediate builds: only if it is a pull request build.
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}

jobs:
test:
runs-on: ${{ matrix.os }}
strategy:
fail-fast: false
matrix:
version:
- "lts"
- "1"
os:
- ubuntu-latest
- macos-latest
- windows-latest
steps:
- uses: actions/checkout@v4
- uses: julia-actions/setup-julia@v2
with:
version: ${{ matrix.version }}
- uses: actions/cache@v4
env:
cache-name: cache-artifacts
with:
path: ~/.julia/artifacts
key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
restore-keys: |
${{ runner.os }}-test-${{ env.cache-name }}-
${{ runner.os }}-test-
${{ runner.os }}-
- name: "Install Dependencies and Run Tests"
run: |
import Pkg
Pkg.Registry.update()
# Install packages present in subdirectories
dev_pks = Pkg.PackageSpec[]
for path in ("lib/SciMLJacobianOperators", "lib/NonlinearSolveBase")
push!(dev_pks, Pkg.PackageSpec(; path))
end
Pkg.develop(dev_pks)
Pkg.instantiate()
Pkg.test(; coverage="user")
shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/NonlinearSolveFirstOrder {0}
- uses: julia-actions/julia-processcoverage@v1
with:
directories: lib/NonlinearSolveFirstOrder/src,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext,lib/SciMLJacobianOperators/src
- uses: codecov/codecov-action@v4
with:
file: lcov.info
token: ${{ secrets.CODECOV_TOKEN }}
verbose: true
fail_ci_if_error: true
49 changes: 49 additions & 0 deletions common/common_nlls_testing.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
using NonlinearSolveBase, SciMLBase, StableRNGs, ForwardDiff, Random, LinearAlgebra

true_function(x, θ) = @. θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4])
true_function(y, x, θ) = (@. y = θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4]))

θ_true = [1.0, 0.1, 2.0, 0.5]
x = [-1.0, -0.5, 0.0, 0.5, 1.0]

const y_target = true_function(x, θ_true)

function loss_function(θ, p)
= true_function(p, θ)
return.- y_target
end

function loss_function(resid, θ, p)
true_function(resid, p, θ)
resid .= resid .- y_target
return resid
end

θ_init = θ_true .+ randn!(StableRNG(0), similar(θ_true)) * 0.1

function vjp(v, θ, p)
resid = zeros(length(p))
J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ)
return vec(v' * J)
end

function vjp!(Jv, v, θ, p)
resid = zeros(length(p))
J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ)
mul!(vec(Jv), transpose(J), v)
return nothing
end

prob_oop = NonlinearLeastSquaresProblem{false}(loss_function, θ_init, x)
prob_iip = NonlinearLeastSquaresProblem{true}(
NonlinearFunction(loss_function; resid_prototype = zero(y_target)), θ_init, x
)
prob_oop_vjp = NonlinearLeastSquaresProblem(
NonlinearFunction{false}(loss_function; vjp), θ_init, x
)
prob_iip_vjp = NonlinearLeastSquaresProblem(
NonlinearFunction{true}(loss_function; resid_prototype = zero(y_target), vjp = vjp!),
θ_init, x
)

export prob_oop, prob_iip, prob_oop_vjp, prob_iip_vjp
File renamed without changes.
4 changes: 3 additions & 1 deletion lib/NonlinearSolveBase/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -83,12 +83,14 @@ julia = "1.10"

[extras]
Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[targets]
test = ["Aqua", "DiffEqBase", "ExplicitImports", "ForwardDiff", "InteractiveUtils", "SparseArrays", "Test"]
test = ["Aqua", "BandedMatrices", "DiffEqBase", "ExplicitImports", "ForwardDiff", "InteractiveUtils", "LinearAlgebra", "SparseArrays", "Test"]
3 changes: 1 addition & 2 deletions lib/NonlinearSolveBase/src/NonlinearSolveBase.jl
Original file line number Diff line number Diff line change
Expand Up @@ -57,8 +57,7 @@ include("solve.jl")
@compat(public, (L2_NORM, Linf_NORM, NAN_CHECK, UNITLESS_ABS2, get_tolerance))
@compat(public, (nonlinearsolve_forwarddiff_solve, nonlinearsolve_dual_solution))
@compat(public,
(select_forward_mode_autodiff, select_reverse_mode_autodiff,
select_jacobian_autodiff))
(select_forward_mode_autodiff, select_reverse_mode_autodiff, select_jacobian_autodiff))

# public for NonlinearSolve.jl and subpackages to use
@compat(public, (InternalAPI, supports_line_search, supports_trust_region, set_du!))
Expand Down
2 changes: 1 addition & 1 deletion lib/NonlinearSolveBase/src/linear_solve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ function (cache::NativeJLLinearSolveCache)(;
b === nothing || (cache.b = b)

if linu !== nothing && ArrayInterface.can_setindex(linu) &&
applicable(ldiv!, linu, cache.A, cache.b)
applicable(ldiv!, linu, cache.A, cache.b) && applicable(ldiv!, cache.A, linu)
ldiv!(linu, cache.A, cache.b)
res = linu
else
Expand Down
6 changes: 6 additions & 0 deletions lib/NonlinearSolveBase/src/utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,12 @@ convert_real(::Type{T}, ::Nothing) where {T} = nothing
convert_real(::Type{T}, x) where {T} = real(T(x))

restructure(::Number, x::Number) = x
function restructure(
y::T1, x::T2
) where {T1 <: AbstractSciMLOperator, T2 <: AbstractSciMLOperator}
@assert size(y) == size(x) "cannot restructure operators. ensure their sizes match."
return x
end
restructure(y, x) = ArrayInterface.restructure(y, x)

function safe_similar(x, args...; kwargs...)
Expand Down
9 changes: 9 additions & 0 deletions lib/NonlinearSolveBase/test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -24,4 +24,13 @@ using InteractiveUtils, Test
@test check_no_stale_explicit_imports(NonlinearSolveBase) === nothing
@test check_all_qualified_accesses_via_owners(NonlinearSolveBase) === nothing
end

@testset "Banded Matrix vcat" begin
using BandedMatrices, LinearAlgebra, SparseArrays

b = BandedMatrix(Ones(5, 5), (1, 1))
d = Diagonal(ones(5, 5))

@test NonlinearSolveBase.Utils.faster_vcat(b, d) == vcat(sparse(b), d)
end
end
25 changes: 22 additions & 3 deletions lib/NonlinearSolveFirstOrder/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@ ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb"
Expand All @@ -27,41 +26,61 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
ADTypes = "1.9.0"
Aqua = "0.8"
ArrayInterface = "7.16.0"
BandedMatrices = "1.7.5"
BenchmarkTools = "1.5.0"
CommonSolve = "0.2.4"
ConcreteStructs = "0.2.3"
DiffEqBase = "6.155.3"
Enzyme = "0.13.12"
ExplicitImports = "1.5"
FiniteDiff = "2.26.0"
ForwardDiff = "0.10.36"
Hwloc = "3"
InteractiveUtils = "<0.0.1, 1"
LineSearch = "0.1.4"
LinearAlgebra = "1.11.0"
LineSearches = "7.3.0"
LinearAlgebra = "1.10"
LinearSolve = "2.36.1"
MaybeInplace = "0.1.4"
NonlinearProblemLibrary = "0.1.2"
NonlinearSolveBase = "1.1"
Pkg = "1.10"
PrecompileTools = "1.2"
Random = "1.10"
ReTestItems = "1.24"
Reexport = "1"
SciMLBase = "2.54"
SciMLJacobianOperators = "0.1.0"
Setfield = "1.1.1"
SparseConnectivityTracer = "0.6.8"
SparseMatrixColorings = "0.4.8"
StableRNGs = "1"
StaticArrays = "1.9.8"
StaticArraysCore = "1.4.3"
Test = "1.10"
Zygote = "0.6.72"
julia = "1.10"

[extras]
Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d"
InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b"
LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

[targets]
test = ["Aqua", "ExplicitImports", "Hwloc", "InteractiveUtils", "NonlinearProblemLibrary", "Pkg", "ReTestItems", "StableRNGs", "Test"]
test = ["Aqua", "BandedMatrices", "BenchmarkTools", "Enzyme", "ExplicitImports", "Hwloc", "InteractiveUtils", "LineSearch", "LineSearches", "NonlinearProblemLibrary", "Pkg", "Random", "ReTestItems", "SparseConnectivityTracer", "SparseMatrixColorings", "StableRNGs", "StaticArrays", "Test", "Zygote"]
3 changes: 2 additions & 1 deletion lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,8 @@ function LevenbergMarquardt(;
autodiff,
vjp_autodiff,
jvp_autodiff,
name = :LevenbergMarquardt
name = :LevenbergMarquardt,
concrete_jac = Val(true)
)
end

Expand Down
4 changes: 1 addition & 3 deletions lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl
Original file line number Diff line number Diff line change
Expand Up @@ -72,9 +72,7 @@ function InternalAPI.init(
) where {F}
T = promote_type(eltype(u), eltype(fu))
return SwitchedEvolutionRelaxationCache(
internalnorm(fu),
T(inv(initial_damping)),
internalnorm
internalnorm(fu), T(inv(initial_damping)), internalnorm
)
end

Expand Down
67 changes: 35 additions & 32 deletions lib/NonlinearSolveFirstOrder/src/solve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -89,53 +89,52 @@ end
kwargs
end

# XXX: Implement
# function __reinit_internal!(
# cache::GeneralizedFirstOrderAlgorithmCache{iip}, args...; p = cache.p, u0 = cache.u,
# alias_u0::Bool = false, maxiters = 1000, maxtime = nothing, kwargs...) where {iip}
# if iip
# recursivecopy!(cache.u, u0)
# cache.prob.f(cache.fu, cache.u, p)
# else
# cache.u = __maybe_unaliased(u0, alias_u0)
# set_fu!(cache, cache.prob.f(cache.u, p))
# end
# cache.p = p

# __reinit_internal!(cache.stats)
# cache.nsteps = 0
# cache.maxiters = maxiters
# cache.maxtime = maxtime
# cache.total_time = 0.0
# cache.force_stop = false
# cache.retcode = ReturnCode.Default
# cache.make_new_jacobian = true

# reset!(cache.trace)
# reinit!(cache.termination_cache, get_fu(cache), get_u(cache); kwargs...)
# reset_timer!(cache.timer)
# end
function InternalAPI.reinit_self!(
cache::GeneralizedFirstOrderAlgorithmCache, args...; p = cache.p, u0 = cache.u,
alias_u0::Bool = false, maxiters = 1000, maxtime = nothing, kwargs...
)
Utils.reinit_common!(cache, u0, p, alias_u0)

InternalAPI.reinit!(cache.stats)
cache.nsteps = 0
cache.maxiters = maxiters
cache.maxtime = maxtime
cache.total_time = 0.0
cache.force_stop = false
cache.retcode = ReturnCode.Default
cache.make_new_jacobian = true

NonlinearSolveBase.reset!(cache.trace)
SciMLBase.reinit!(
cache.termination_cache, NonlinearSolveBase.get_fu(cache),
NonlinearSolveBase.get_u(cache); kwargs...
)
NonlinearSolveBase.reset_timer!(cache.timer)
return
end

NonlinearSolveBase.@internal_caches(GeneralizedFirstOrderAlgorithmCache,
:jac_cache, :descent_cache, :linesearch_cache, :trustregion_cache)

function SciMLBase.__init(
prob::AbstractNonlinearProblem, alg::GeneralizedFirstOrderAlgorithm,
args...; stats = NLStats(0, 0, 0, 0, 0), alias_u0 = false, maxiters = 1000,
prob::AbstractNonlinearProblem, alg::GeneralizedFirstOrderAlgorithm, args...;
stats = NLStats(0, 0, 0, 0, 0), alias_u0 = false, maxiters = 1000,
abstol = nothing, reltol = nothing, maxtime = nothing,
termination_condition = nothing, internalnorm = L2_NORM,
linsolve_kwargs = (;), kwargs...
)
@set! alg.autodiff = NonlinearSolveBase.select_jacobian_autodiff(prob, alg.autodiff)
@set! alg.jvp_autodiff = if alg.jvp_autodiff === nothing && alg.autodiff !== nothing &&
provided_jvp_autodiff = alg.jvp_autodiff !== nothing
@set! alg.jvp_autodiff = if !provided_jvp_autodiff && alg.autodiff !== nothing &&
(ADTypes.mode(alg.autodiff) isa ADTypes.ForwardMode ||
ADTypes.mode(alg.autodiff) isa
ADTypes.ForwardOrReverseMode)
NonlinearSolveBase.select_forward_mode_autodiff(prob, alg.autodiff)
else
NonlinearSolveBase.select_forward_mode_autodiff(prob, alg.jvp_autodiff)
end
@set! alg.vjp_autodiff = if alg.vjp_autodiff === nothing && alg.autodiff !== nothing &&
provided_vjp_autodiff = alg.vjp_autodiff !== nothing
@set! alg.vjp_autodiff = if !provided_vjp_autodiff && alg.autodiff !== nothing &&
(ADTypes.mode(alg.autodiff) isa ADTypes.ReverseMode ||
ADTypes.mode(alg.autodiff) isa
ADTypes.ForwardOrReverseMode)
Expand Down Expand Up @@ -185,7 +184,7 @@ function SciMLBase.__init(
error("Trust Region not supported by $(alg.descent).")
trustregion_cache = InternalAPI.init(
prob, alg.trustregion, prob.f, fu, u, prob.p;
stats, internalnorm, kwargs...
alg.vjp_autodiff, alg.jvp_autodiff, stats, internalnorm, kwargs...
)
globalization = Val(:TrustRegion)
end
Expand All @@ -194,7 +193,11 @@ function SciMLBase.__init(
NonlinearSolveBase.supports_line_search(alg.descent) ||
error("Line Search not supported by $(alg.descent).")
linesearch_cache = CommonSolve.init(
prob, alg.linesearch, fu, u; stats, internalnorm, kwargs...
prob, alg.linesearch, fu, u; stats, internalnorm,
autodiff = ifelse(
provided_jvp_autodiff, alg.jvp_autodiff, alg.vjp_autodiff
),
kwargs...
)
globalization = Val(:LineSearch)
end
Expand Down
Loading

0 comments on commit d0389b7

Please sign in to comment.