Skip to content

Commit a2ad001

Browse files
authored
Make Krylov solver verbose and public (#157)
1 parent 7c0aa99 commit a2ad001

File tree

2 files changed

+32
-25
lines changed

2 files changed

+32
-25
lines changed

src/ImplicitDifferentiation.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,6 @@ using LinearAlgebra: axpby!, factorize, lu
2323
include("implicit_function.jl")
2424
include("operators.jl")
2525

26-
export ImplicitFunction
26+
export ImplicitFunction, KrylovLinearSolver
2727

2828
end

src/implicit_function.jl

Lines changed: 31 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,43 @@
11
"""
22
KrylovLinearSolver
33
4-
Callable object that can solve linear systems `As = b` and `AS = b` in the same way as the built-in `\\`.
4+
Callable object that can solve linear systems `Ax = b` and `AX = B` in the same way as the built-in `\\`.
55
Uses an iterative solver from [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl) under the hood.
66
7-
# Note
7+
# Constructor
88
9-
This name is not exported, and thus not part of the public API, but it is used in the [`ImplicitFunction`](@ref) constructors.
10-
"""
11-
struct KrylovLinearSolver end
9+
KrylovLinearSolver(; verbose=true)
10+
11+
If `verbose` is `true`, the solver logs a warning in case of failure.
12+
Otherwise it will fail silently, and may return solutions that do not exactly satisfy the linear system.
13+
14+
# Callable behavior
1215
13-
"""
1416
(::KylovLinearSolver)(A, b::AbstractVector)
1517
1618
Solve a linear system with a single right-hand side.
17-
"""
18-
function (::KrylovLinearSolver)(A, b::AbstractVector)
19-
x, stats = gmres(A, b)
20-
return x
21-
end
2219
23-
"""
2420
(::KrylovLinearSolver)(A, B::AbstractMatrix)
2521
2622
Solve a linear system with multiple right-hand sides.
2723
"""
28-
function (::KrylovLinearSolver)(A, B::AbstractMatrix)
24+
Base.@kwdef struct KrylovLinearSolver
25+
verbose::Bool = true
26+
end
27+
28+
function (solver::KrylovLinearSolver)(A, b::AbstractVector)
29+
x, stats = gmres(A, b)
30+
if !stats.solved || stats.inconsistent
31+
solver.verbose &&
32+
@warn "Failed to solve the linear system in the implicit function theorem with `Krylov.gmres`" stats
33+
end
34+
return x
35+
end
36+
37+
function (solver::KrylovLinearSolver)(A, B::AbstractMatrix)
2938
# X, stats = block_gmres(A, B) # https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/854
3039
X = mapreduce(hcat, eachcol(B)) do b
31-
first(gmres(A, b))
40+
solver(A, b)
3241
end
3342
return X
3443
end
@@ -80,6 +89,14 @@ Picks the `lazy` parameter automatically based on the `linear_solver`, using the
8089
8190
Picks the `linear_solver` automatically based on the `lazy` parameter.
8291
92+
# Callable behavior
93+
94+
(implicit::ImplicitFunction)(x::AbstractVector, args...; kwargs...)
95+
96+
Return `implicit.forward(x, args...; kwargs...)`, which can be either an `AbstractVector` `y` or a tuple `(y, z)`.
97+
98+
This call makes `y` differentiable with respect to `x`.
99+
83100
# Function signatures
84101
85102
There are two possible signatures for `forward` and `conditions`, which must be consistent with one another:
@@ -122,9 +139,6 @@ struct ImplicitFunction{
122139
conditions_y_backend::B2
123140
end
124141

125-
"""
126-
127-
"""
128142
function ImplicitFunction{lazy}(
129143
forward::F,
130144
conditions::C;
@@ -163,13 +177,6 @@ function Base.show(io::IO, implicit::ImplicitFunction{lazy}) where {lazy}
163177
)
164178
end
165179

166-
"""
167-
(implicit::ImplicitFunction)(x::AbstractVector, args...; kwargs...)
168-
169-
Return `implicit.forward(x, args...; kwargs...)`, which can be either an `AbstractVector` `y` or a tuple `(y, z)`.
170-
171-
This call makes `y` differentiable with respect to `x`.
172-
"""
173180
function (implicit::ImplicitFunction)(x::AbstractVector, args...; kwargs...)
174181
y_or_yz = implicit.forward(x, args...; kwargs...)
175182
return y_or_yz

0 commit comments

Comments
 (0)