Skip to content

Commit

Permalink
materialize the multi-step scheme
Browse files Browse the repository at this point in the history
  • Loading branch information
avik-pal committed Feb 13, 2024
1 parent fb88484 commit e22e052
Show file tree
Hide file tree
Showing 32 changed files with 291 additions and 155 deletions.
4 changes: 3 additions & 1 deletion Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ version = "3.6.0"

[deps]
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
Expand Down Expand Up @@ -55,12 +56,13 @@ NonlinearSolveZygoteExt = "Zygote"

[compat]
ADTypes = "0.2.6"
Accessors = "0.1"
Aqua = "0.8"
ArrayInterface = "7.7"
BandedMatrices = "1.4"
BenchmarkTools = "1.4"
ConcreteStructs = "0.2.3"
CUDA = "5.1"
ConcreteStructs = "0.2.3"
DiffEqBase = "6.146.0"
Enzyme = "0.11.11"
FastBroadcast = "0.2.8"
Expand Down
2 changes: 1 addition & 1 deletion docs/make.jl
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
using Documenter, DocumenterCitations
using NonlinearSolve,
SimpleNonlinearSolve, Sundials, SteadyStateDiffEq, SciMLBase, DiffEqBase
SimpleNonlinearSolve, Sundials, SteadyStateDiffEq, SciMLBase, DiffEqBase

cp(joinpath(@__DIR__, "Manifest.toml"), joinpath(@__DIR__, "src/assets/Manifest.toml"),
force = true)
Expand Down
2 changes: 1 addition & 1 deletion docs/pages.jl
Original file line number Diff line number Diff line change
Expand Up @@ -43,5 +43,5 @@ pages = [
"devdocs/operators.md",
"devdocs/algorithm_helpers.md"],
"Release Notes" => "release_notes.md",
"References" => "references.md",
"References" => "references.md"
]
16 changes: 10 additions & 6 deletions docs/src/basics/faq.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,19 @@ myfun(x, lv) = x * sin(x) - lv
function f(out, levels, u0)
for i in 1:N
out[i] = solve(IntervalNonlinearProblem{false}(IntervalNonlinearFunction{false}(myfun),
u0, levels[i]), Falsi()).u
out[i] = solve(
IntervalNonlinearProblem{false}(IntervalNonlinearFunction{false}(myfun),
u0, levels[i]),
Falsi()).u
end
end
function f2(out, levels, u0)
for i in 1:N
out[i] = solve(NonlinearProblem{false}(NonlinearFunction{false}(myfun),
u0, levels[i]), SimpleNewtonRaphson()).u
out[i] = solve(
NonlinearProblem{false}(NonlinearFunction{false}(myfun),
u0, levels[i]),
SimpleNewtonRaphson()).u
end
end
Expand Down Expand Up @@ -68,7 +72,7 @@ differentiate the function based on the input types. However, this function has
`xx = [1.0, 2.0, 3.0, 4.0]` followed by a `xx[1] = var[1] - v_true[1]` where `var` might
be a Dual number. This causes the error. To fix it:

1. Specify the `autodiff` to be `AutoFiniteDiff`
1. Specify the `autodiff` to be `AutoFiniteDiff`

```@example dual_error_faq
sol = solve(prob_oop, LevenbergMarquardt(; autodiff = AutoFiniteDiff()); maxiters = 10000,
Expand All @@ -77,7 +81,7 @@ sol = solve(prob_oop, LevenbergMarquardt(; autodiff = AutoFiniteDiff()); maxiter

This worked but, Finite Differencing is not the recommended approach in any scenario.

2. Rewrite the function to use
2. Rewrite the function to use
[PreallocationTools.jl](https://github.com/SciML/PreallocationTools.jl) or write it as

```@example dual_error_faq
Expand Down
13 changes: 8 additions & 5 deletions docs/src/basics/sparsity_detection.md
Original file line number Diff line number Diff line change
Expand Up @@ -22,17 +22,19 @@ NonlinearSolve will automatically perform matrix coloring and use sparse differe
Now you can help the solver further by providing the color vector. This can be done by

```julia
prob = NonlinearProblem(NonlinearFunction(nlfunc; sparsity = jac_prototype,
prob = NonlinearProblem(
NonlinearFunction(nlfunc; sparsity = jac_prototype,
colorvec = colorvec), x0)
# OR
prob = NonlinearProblem(NonlinearFunction(nlfunc; jac_prototype = jac_prototype,
prob = NonlinearProblem(
NonlinearFunction(nlfunc; jac_prototype = jac_prototype,
colorvec = colorvec), x0)
```

If the `colorvec` is not provided, then it is computed on demand.

!!! note

One thing to be careful about in this case is that `colorvec` is dependent on the
autodiff backend used. Forward Mode and Finite Differencing will assume that the
colorvec is the column colorvec, while Reverse Mode will assume that the colorvec is the
Expand All @@ -47,7 +49,8 @@ algorithm you want to use, then you can create your problem as follows:
prob = NonlinearProblem(NonlinearFunction(nlfunc; sparsity = SymbolicsSparsityDetection()),
x0) # Remember to have Symbolics.jl loaded
# OR
prob = NonlinearProblem(NonlinearFunction(nlfunc; sparsity = ApproximateJacobianSparsity()),
prob = NonlinearProblem(
NonlinearFunction(nlfunc; sparsity = ApproximateJacobianSparsity()),
x0)
```

Expand All @@ -73,7 +76,7 @@ loaded, we default to using `SymbolicsSparsityDetection()`, else we default to u
options if those are provided.

!!! warning

If you provide a non-sparse AD, and provide a `sparsity` or `jac_prototype` then
we will use dense AD. This is because, if you provide a specific AD type, we assume
that you know what you are doing and want to override the default choice of `nothing`.
30 changes: 17 additions & 13 deletions docs/src/tutorials/large_systems.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,15 @@

This tutorial is for getting into the extra features of using NonlinearSolve.jl. Solving
ill-conditioned nonlinear systems requires specializing the linear solver on properties of
the Jacobian in order to cut down on the ``\mathcal{O}(n^3)`` linear solve and the
``\mathcal{O}(n^2)`` back-solves. This tutorial is designed to explain the advanced usage of
the Jacobian in order to cut down on the `\mathcal{O}(n^3)` linear solve and the
`\mathcal{O}(n^2)` back-solves. This tutorial is designed to explain the advanced usage of
NonlinearSolve.jl by solving the steady state stiff Brusselator partial differential
equation (BRUSS) using NonlinearSolve.jl.

## Definition of the Brusselator Equation

!!! note

Feel free to skip this section: it simply defines the example problem.

The Brusselator PDE is defined as follows:
Expand Down Expand Up @@ -118,11 +118,11 @@ However, if you know the sparsity of your problem, then you can pass a different
type. For example, a `SparseMatrixCSC` will give a sparse matrix. Other sparse matrix types
include:

- Bidiagonal
- Tridiagonal
- SymTridiagonal
- BandedMatrix ([BandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BandedMatrices.jl))
- BlockBandedMatrix ([BlockBandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BlockBandedMatrices.jl))
- Bidiagonal
- Tridiagonal
- SymTridiagonal
- BandedMatrix ([BandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BandedMatrices.jl))
- BlockBandedMatrix ([BlockBandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BlockBandedMatrices.jl))

## Approximate Sparsity Detection & Sparse Jacobians

Expand Down Expand Up @@ -213,7 +213,7 @@ choices, see the
`linsolve` choices are any valid [LinearSolve.jl](https://linearsolve.sciml.ai/dev/) solver.

!!! note

Switching to a Krylov linear solver will automatically change the nonlinear problem
solver into Jacobian-free mode, dramatically reducing the memory required. This can be
overridden by adding `concrete_jac=true` to the algorithm.
Expand Down Expand Up @@ -308,10 +308,14 @@ for the exact sparsity detection case, we left out the time it takes to perform
sparsity detection. Let's compare the two by setting the sparsity detection algorithms.

```@example ill_conditioned_nlprob
prob_brusselator_2d_exact = NonlinearProblem(NonlinearFunction(brusselator_2d_loop;
sparsity = SymbolicsSparsityDetection()), u0, p; abstol = 1e-10, reltol = 1e-10)
prob_brusselator_2d_approx = NonlinearProblem(NonlinearFunction(brusselator_2d_loop;
sparsity = ApproximateJacobianSparsity()), u0, p; abstol = 1e-10, reltol = 1e-10)
prob_brusselator_2d_exact = NonlinearProblem(
NonlinearFunction(brusselator_2d_loop;
sparsity = SymbolicsSparsityDetection()),
u0, p; abstol = 1e-10, reltol = 1e-10)
prob_brusselator_2d_approx = NonlinearProblem(
NonlinearFunction(brusselator_2d_loop;
sparsity = ApproximateJacobianSparsity()),
u0, p; abstol = 1e-10, reltol = 1e-10)
@btime solve(prob_brusselator_2d_exact, NewtonRaphson());
@btime solve(prob_brusselator_2d_approx, NewtonRaphson());
Expand Down
6 changes: 3 additions & 3 deletions docs/src/tutorials/modelingtoolkit.md
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,8 @@ u0 = [x => 1.0,
z => 0.0]
ps = [σ => 10.0
ρ => 26.0
β => 8 / 3]
ρ => 26.0
β => 8 / 3]
prob = NonlinearProblem(ns, u0, ps)
sol = solve(prob, NewtonRaphson())
Expand Down Expand Up @@ -65,7 +65,7 @@ eqs = [
0 ~ u2 - cos(u1),
0 ~ u3 - hypot(u1, u2),
0 ~ u4 - hypot(u2, u3),
0 ~ u5 - hypot(u4, u1),
0 ~ u5 - hypot(u4, u1)
]
@named sys = NonlinearSystem(eqs, [u1, u2, u3, u4, u5], [])
```
Expand Down
3 changes: 2 additions & 1 deletion ext/NonlinearSolveMINPACKExt.jl
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@ using MINPACK, NonlinearSolve, SciMLBase
import FastClosures: @closure

function SciMLBase.__solve(prob::Union{NonlinearLeastSquaresProblem,
NonlinearProblem}, alg::CMINPACK, args...; abstol = nothing, maxiters = 1000,
NonlinearProblem},
alg::CMINPACK, args...; abstol = nothing, maxiters = 1000,
alias_u0::Bool = false, show_trace::Val{ShT} = Val(false),
store_trace::Val{StT} = Val(false), termination_condition = nothing,
kwargs...) where {ShT, StT}
Expand Down
51 changes: 30 additions & 21 deletions src/NonlinearSolve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -8,22 +8,23 @@ import Reexport: @reexport
import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_workload

@recompile_invalidations begin
using ADTypes, ConcreteStructs, DiffEqBase, FastBroadcast, FastClosures, LazyArrays,
LineSearches, LinearAlgebra, LinearSolve, MaybeInplace, Preferences, Printf,
SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools
using Accessors, ADTypes, ConcreteStructs, DiffEqBase, FastBroadcast, FastClosures,
LazyArrays, LineSearches, LinearAlgebra, LinearSolve, MaybeInplace, Preferences,
Printf, SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools

import ArrayInterface: undefmatrix, can_setindex, restructure, fast_scalar_indexing
import DiffEqBase: AbstractNonlinearTerminationMode,
AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode,
NonlinearSafeTerminationReturnCode, get_termination_mode
AbstractSafeNonlinearTerminationMode,
AbstractSafeBestNonlinearTerminationMode,
NonlinearSafeTerminationReturnCode, get_termination_mode
import FiniteDiff
import ForwardDiff
import ForwardDiff: Dual
import LinearSolve: ComposePreconditioner, InvPreconditioner, needs_concrete_A
import RecursiveArrayTools: recursivecopy!, recursivefill!

import SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem,
AbstractSciMLOperator, NLStats, _unwrap_val, has_jac, isinplace
AbstractSciMLOperator, NLStats, _unwrap_val, has_jac, isinplace
import SparseDiffTools: AbstractSparsityDetection, AutoSparseEnzyme
import StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix, MMatrix
end
Expand Down Expand Up @@ -98,20 +99,28 @@ include("default.jl")
probs_nlls = NonlinearLeastSquaresProblem[]
nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), [0.1, 0.0]),
(NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), [0.1, 0.1]),
(NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p,
resid_prototype = zeros(1)), [0.1, 0.0]),
(NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p),
resid_prototype = zeros(4)), [0.1, 0.1]))
(
NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p,
resid_prototype = zeros(1)),
[0.1, 0.0]),
(
NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p),
resid_prototype = zeros(4)),
[0.1, 0.1]))
for (fn, u0) in nlfuncs
push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0))
end
nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), Float32[0.1, 0.0]),
(NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)),
Float32[0.1, 0.1]),
(NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p,
resid_prototype = zeros(Float32, 1)), Float32[0.1, 0.0]),
(NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p),
resid_prototype = zeros(Float32, 4)), Float32[0.1, 0.1]))
(
NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p,
resid_prototype = zeros(Float32, 1)),
Float32[0.1, 0.0]),
(
NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p),
resid_prototype = zeros(Float32, 4)),
Float32[0.1, 0.1]))
for (fn, u0) in nlfuncs
push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0f0))
end
Expand All @@ -133,21 +142,21 @@ end

# Core Algorithms
export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden, DFSane,
MultiStepNonlinearSolver
MultiStepNonlinearSolver
export GaussNewton, LevenbergMarquardt, TrustRegion
export NonlinearSolvePolyAlgorithm,
RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg
RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg

# Extension Algorithms
export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL,
FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL
FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL

# Advanced Algorithms -- Without Bells and Whistles
export GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm, GeneralizedDFSane

# Descent Algorithms
export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent,
GeodesicAcceleration, GenericMultiStepDescent
GeodesicAcceleration, GenericMultiStepDescent
## Multistep Algorithms
export MultiStepSchemes

Expand All @@ -159,9 +168,9 @@ export RadiusUpdateSchemes

# Export the termination conditions from DiffEqBase
export SteadyStateDiffEqTerminationMode, SimpleNonlinearSolveTerminationMode,
NormTerminationMode, RelTerminationMode, RelNormTerminationMode, AbsTerminationMode,
AbsNormTerminationMode, RelSafeTerminationMode, AbsSafeTerminationMode,
RelSafeBestTerminationMode, AbsSafeBestTerminationMode
NormTerminationMode, RelTerminationMode, RelNormTerminationMode, AbsTerminationMode,
AbsNormTerminationMode, RelSafeTerminationMode, AbsSafeTerminationMode,
RelSafeBestTerminationMode, AbsSafeBestTerminationMode

# Tracing Functionality
export TraceAll, TraceMinimal, TraceWithJacobianConditionNumber
Expand Down
6 changes: 4 additions & 2 deletions src/abstract_types.jl
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@ squares solver.
### `__internal_init` specification
```julia
__internal_init(prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u;
__internal_init(
prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u;
pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), abstol = nothing,
reltol = nothing, alias_J::Bool = true, shared::Val{N} = Val(1),
kwargs...) where {INV, N, uType, iip} --> AbstractDescentCache
Expand Down Expand Up @@ -225,7 +226,8 @@ Abstract Type for Damping Functions in DampedNewton.
### `__internal_init` specification
```julia
__internal_init(prob::AbstractNonlinearProblem, f::AbstractDampingFunction, initial_damping,
__internal_init(
prob::AbstractNonlinearProblem, f::AbstractDampingFunction, initial_damping,
J, fu, u, args...; internal_norm = DEFAULT_NORM,
kwargs...) --> AbstractDampingFunctionCache
```
Expand Down
3 changes: 2 additions & 1 deletion src/algorithms/broyden.jl
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,8 @@ search.
useful for specific problems, but whether it will work may depend strongly on the
problem
"""
function Broyden(; max_resets = 100, linesearch = NoLineSearch(), reset_tolerance = nothing,
function Broyden(;
max_resets = 100, linesearch = NoLineSearch(), reset_tolerance = nothing,
init_jacobian::Val{IJ} = Val(:identity), autodiff = nothing, alpha = nothing,
update_rule::Val{UR} = Val(:good_broyden)) where {IJ, UR}
if IJ === :identity
Expand Down
3 changes: 2 additions & 1 deletion src/algorithms/dfsane.jl
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@ For other keyword arguments, see [`RobustNonMonotoneLineSearch`](@ref).
function DFSane(; σ_min = 1 // 10^10, σ_max = 1e10, σ_1 = 1, M::Int = 10, γ = 1 // 10^4,
τ_min = 1 // 10, τ_max = 1 // 2, n_exp::Int = 2, max_inner_iterations::Int = 100,
η_strategy::ETA = (fn_1, n, x_n, f_n) -> fn_1 / n^2) where {ETA}
linesearch = RobustNonMonotoneLineSearch(; gamma = γ, sigma_1 = σ_1, M, tau_min = τ_min,
linesearch = RobustNonMonotoneLineSearch(;
gamma = γ, sigma_1 = σ_1, M, tau_min = τ_min,
tau_max = τ_max, n_exp, η_strategy, maxiters = max_inner_iterations)
return GeneralizedDFSane{:DFSane}(linesearch, σ_min, σ_max, nothing)
end
Loading

0 comments on commit e22e052

Please sign in to comment.