14544069614

Committed 19 Apr 2025 12:57AM UTC coverage: 94.777% (-0.009%) from 94.786%

Build # 14544069614

Build Type

push

github

Committed by

web-flow

Commit Message

Better exact diagonalisation (#306)

# Changes
* Implement a `LinearMap` from
[LinearMaps.jl](https://github.com/JuliaLinearAlgebra/LinearMaps.jl),
which allows us to use operators as matrices without storing the
matrices in memory.
* Use the `LinearMap` to implement matrix-free variants of
`KrylovKitSolver`, `ArpackSolver`, and `LOBPCGSolver`.
* Refactor of the `ExactDiagonalizationProblem` and its surrounding
infrastructure.

# Changed behaviour
* `ArpackSolver` and `LOBPCGSolver` now use the matrix-free algorithm by
default.

Run Details

208 of 209 new or added lines in 9 files covered. (99.52%)

19 existing lines in 3 files now uncovered.

6932 of 7314 relevant lines covered (94.78%)

11143439.58 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

98.53

/src/ExactDiagonalization/init_and_solvers.jl

"""
    init(p::ExactDiagonalizationProblem, [algorithm]; kwargs...)

Initialize a solver for an [`ExactDiagonalizationProblem`](@ref) `p` with an optional
`algorithm`. Returns a solver instance that can be solved with
[`solve`](@ref solve(::ExactDiagonalizationProblem)).

For a description of the keyword arguments, see the documentation for
[`ExactDiagonalizationProblem`](@ref).
"""
function CommonSolve.init( # no algorithm specified as positional argument
    prob::ExactDiagonalizationProblem;
    kwargs...
)
    kwargs = (; prob.kwargs..., kwargs...) # remove duplicates
    algorithm = get(kwargs, :algorithm, LinearAlgebraSolver())
    kwargs = delete(kwargs, :algorithm)
    new_prob = ExactDiagonalizationProblem(prob.hamiltonian, prob.initial_vector; kwargs...)
    return init(new_prob, algorithm)
end

# TODO since this is the same for all solvers, maybe move it to ExactDiagonalizationProblem?
function _set_up_starting_address(v0, ham)
    if isnothing(v0)
        addr_or_vec = starting_address(ham)
    elseif allows_address_type(ham, v0) ||
            v0 isa Union{NTuple,Vector} && allows_address_type(ham, eltype(v0))
        addr_or_vec = v0
    elseif v0 isa FrozenDVec
        addr_or_vec = keys(v0)
    else
        throw(ArgumentError("Invalid starting vector in `ExactDiagonalizationProblem`."))
    end

    return addr_or_vec
end
function _set_up_initial_vector(ham, v0, basis)
    T = float(eltype(ham))
    if isnothing(v0)
        return rand(T, length(basis))
    end

    if v0 isa Union{NTuple, AbstractVector} && eltype(v0) == eltype(basis)
        v0_dvec = Dict(addr => 1.0 for addr in v0)
    elseif v0 isa eltype(basis)
        v0_dvec = Dict(v0 => 1.0)
    elseif v0 isa FrozenDVec
        v0_dvec = Dict(pairs(v0))
    else
        @assert false # this should be unreachable
    end

    return [T(get(v0_dvec, b, zero(valtype(v0_dvec)))) for b in basis]
end

struct IterativeEDSolver{A,P,LM,T<:Number,F}
    algorithm::A
    problem::P
    linear_map::LM
    initial_vector::Vector{T}
    basis::Vector{F}
    solver_kwargs::NamedTuple
end
function Base.show(io::IO, s::IterativeEDSolver)
    io = IOContext(io, :compact => true)
    print(io, "IterativeEDSolver\n with algorithm $(s.algorithm) for hamiltonian = ")
    show(io, s.problem.hamiltonian)
    print(io, ",\n  kwargs = ")
    show(io, s.solver_kwargs)
    print(io, "\n)")
end

function CommonSolve.init(
    prob::ExactDiagonalizationProblem, algorithm::AbstractAlgorithm{true}; kwargs...
)
    !ishermitian(prob.hamiltonian) && algorithm isa LOBPCGSolver &&
        throw(ArgumentError("LOBPCGSolver() is not suitable for non-hermitian matrices."))

    # Merge keyword arguments from problem, algorithm and ones passed to this function
    # and split them into sets that are passed to LinearMap and ones that
    # are left for the solver.

    kwargs = (; prob.kwargs..., algorithm.kwargs..., kwargs...)
    linmap_kwargs, solver_kwargs = split_keys(kwargs, :basis, :full_basis)

    # determine the starting address or vector and seed address to build the matrix from
    addr_or_vec = _set_up_starting_address(
        prob.initial_vector, prob.hamiltonian
    )

    # create the LinearMap
    linmap = LinearMap(prob.hamiltonian; starting_address=addr_or_vec, linmap_kwargs...)
    basis = linmap.basis

    initial_vector = _set_up_initial_vector(prob.hamiltonian, prob.initial_vector, basis)

    return IterativeEDSolver(algorithm, prob, linmap, initial_vector, basis, solver_kwargs)
end
function CommonSolve.init(
    prob::ExactDiagonalizationProblem, algorithm::AbstractAlgorithm{false}; kwargs...
)
    !ishermitian(prob.hamiltonian) && algorithm isa LOBPCGSolver &&
        throw(ArgumentError("LOBPCGSolver() is not suitable for non-hermitian matrices."))

    # Merge keyword arguments from problem, algorithm and ones passed to this function
    # and split them into sets that are passed to BasisSetRepresentation and ones that
    # are left for the solver.
    kwargs = (; prob.kwargs..., algorithm.kwargs..., kwargs...)
    bsr_kwargs, solver_kwargs = split_keys(
        kwargs,
        :sizelim, :cutoff, :filter, :nnzs, :col_hint, :sort, :max_depth, :minimum_size
    )

    # determine the starting address or vector and seed address to build the matrix from
    addr_or_vec = _set_up_starting_address(
        prob.initial_vector, prob.hamiltonian
    )

    # create the BasisSetRepresentation
    bsr = BasisSetRepresentation(prob.hamiltonian, addr_or_vec; bsr_kwargs...)
    matrix = bsr.sparse_matrix
    basis = bsr.basis

    initial_vector = _set_up_initial_vector(prob.hamiltonian, prob.initial_vector, basis)

    return IterativeEDSolver(algorithm, prob, matrix, initial_vector, basis, solver_kwargs)
end

struct DenseEDSolver{P,A,BSR}
    problem::P
    algorithm::A
    basis_set_rep::BSR
    solver_kwargs::NamedTuple
end
function Base.show(io::IO, s::DenseEDSolver)
    io = IOContext(io, :compact => true)
    print(io, "DenseEDSolver\n for hamiltonian = ")
    show(io, s.problem.hamiltonian)
    print(io, ",\n  kwargs = ")
    show(io, s.solver_kwargs)
    print(io, "\n)")
end

function CommonSolve.init(
    prob::ExactDiagonalizationProblem, algorithm::LinearAlgebraSolver; kwargs...
)
    # Merge keyword arguments from problem, algorithm and ones passed to this function
    # and split them into sets that are passed to BasisSetRepresentation and ones that
    # are left for the solver.
    kwargs = (; sizelim=1e5, prob.kwargs..., algorithm.kwargs..., kwargs...)
    bsr_kwargs, solver_kwargs = split_keys(
        kwargs,
        :sizelim, :cutoff, :filter, :nnzs, :col_hint, :sort, :max_depth, :minimum_size
    )

    # determine the seed address to build the matrix from
    addr_or_vec = _set_up_starting_address(
        prob.initial_vector, prob.hamiltonian
    )

    # create the BasisSetRepresentation
    bsr = BasisSetRepresentation(prob.hamiltonian, addr_or_vec; bsr_kwargs...)

    return DenseEDSolver(prob, algorithm, bsr, solver_kwargs)
end

1	"""
2	init(p::ExactDiagonalizationProblem, [algorithm]; kwargs...)
3
4	Initialize a solver for an [`ExactDiagonalizationProblem`](@ref) `p` with an optional
5	`algorithm`. Returns a solver instance that can be solved with
6	[`solve`](@ref solve(::ExactDiagonalizationProblem)).
7
8	For a description of the keyword arguments, see the documentation for
9	[`ExactDiagonalizationProblem`](@ref).
10	"""
11	function CommonSolve.init( # no algorithm specified as positional argument	86✔
12	prob::ExactDiagonalizationProblem;
13	kwargs...
14	)
15	kwargs = (; prob.kwargs..., kwargs...) # remove duplicates	55✔
16	algorithm = get(kwargs, :algorithm, LinearAlgebraSolver())	43✔
17	kwargs = delete(kwargs, :algorithm)	43✔
18	new_prob = ExactDiagonalizationProblem(prob.hamiltonian, prob.initial_vector; kwargs...)	43✔
19	return init(new_prob, algorithm)	43✔
20	end
21
22	# TODO since this is the same for all solvers, maybe move it to ExactDiagonalizationProblem?
23	function _set_up_starting_address(v0, ham)	304✔
24	if isnothing(v0)	304✔
25	addr_or_vec = starting_address(ham)	158✔
26	elseif allows_address_type(ham, v0) \|\|	146✔
27	v0 isa Union{NTuple,Vector} && allows_address_type(ham, eltype(v0))
28	addr_or_vec = v0	72✔
29	elseif v0 isa FrozenDVec	74✔
30	addr_or_vec = keys(v0)	72✔
31	else
32	throw(ArgumentError("Invalid starting vector in `ExactDiagonalizationProblem`."))	2✔
33	end
34
35	return addr_or_vec	302✔
36	end
37	function _set_up_initial_vector(ham, v0, basis)	220✔
38	T = float(eltype(ham))	220✔
39	if isnothing(v0)	220✔
40	return rand(T, length(basis))	799✔
41	end
42
43	if v0 isa Union{NTuple, AbstractVector} && eltype(v0) == eltype(basis)	120✔
44	v0_dvec = Dict(addr => 1.0 for addr in v0)	60✔
45	elseif v0 isa eltype(basis)	80✔
46	v0_dvec = Dict(v0 => 1.0)	40✔
47	elseif v0 isa FrozenDVec	60✔
48	v0_dvec = Dict(pairs(v0))	60✔
49	else
NEW 50	@assert false # this should be unreachable	×
51	end
52
53	return [T(get(v0_dvec, b, zero(valtype(v0_dvec)))) for b in basis]	120✔
54	end
55
56	struct IterativeEDSolver{A,P,LM,T<:Number,F}
57	algorithm::A	220✔
58	problem::P
59	linear_map::LM
60	initial_vector::Vector{T}
61	basis::Vector{F}
62	solver_kwargs::NamedTuple
63	end
64	function Base.show(io::IO, s::IterativeEDSolver)	21✔
65	io = IOContext(io, :compact => true)	21✔
66	print(io, "IterativeEDSolver\n with algorithm $(s.algorithm) for hamiltonian = ")	21✔
67	show(io, s.problem.hamiltonian)	21✔
68	print(io, ",\n kwargs = ")	21✔
69	show(io, s.solver_kwargs)	21✔
70	print(io, "\n)")	21✔
71	end
72
73	function CommonSolve.init(	266✔
74	prob::ExactDiagonalizationProblem, algorithm::AbstractAlgorithm{true}; kwargs...
75	)
76	!ishermitian(prob.hamiltonian) && algorithm isa LOBPCGSolver &&	133✔
77	throw(ArgumentError("LOBPCGSolver() is not suitable for non-hermitian matrices."))
78
79	# Merge keyword arguments from problem, algorithm and ones passed to this function
80	# and split them into sets that are passed to LinearMap and ones that
81	# are left for the solver.
82
83	kwargs = (; prob.kwargs..., algorithm.kwargs..., kwargs...)	228✔
84	linmap_kwargs, solver_kwargs = split_keys(kwargs, :basis, :full_basis)	131✔
85
86	# determine the starting address or vector and seed address to build the matrix from
87	addr_or_vec = _set_up_starting_address(	131✔
88	prob.initial_vector, prob.hamiltonian
89	)
90
91	# create the LinearMap
92	linmap = LinearMap(prob.hamiltonian; starting_address=addr_or_vec, linmap_kwargs...)	129✔
93	basis = linmap.basis	129✔
94
95	initial_vector = _set_up_initial_vector(prob.hamiltonian, prob.initial_vector, basis)	682✔
96
97	return IterativeEDSolver(algorithm, prob, linmap, initial_vector, basis, solver_kwargs)	129✔
98	end
99	function CommonSolve.init(	186✔
100	prob::ExactDiagonalizationProblem, algorithm::AbstractAlgorithm{false}; kwargs...
101	)
102	!ishermitian(prob.hamiltonian) && algorithm isa LOBPCGSolver &&	93✔
103	throw(ArgumentError("LOBPCGSolver() is not suitable for non-hermitian matrices."))
104
105	# Merge keyword arguments from problem, algorithm and ones passed to this function
106	# and split them into sets that are passed to BasisSetRepresentation and ones that
107	# are left for the solver.
108	kwargs = (; prob.kwargs..., algorithm.kwargs..., kwargs...)	182✔
109	bsr_kwargs, solver_kwargs = split_keys(	91✔
110	kwargs,
111	:sizelim, :cutoff, :filter, :nnzs, :col_hint, :sort, :max_depth, :minimum_size
112	)
113
114	# determine the starting address or vector and seed address to build the matrix from
115	addr_or_vec = _set_up_starting_address(	91✔
116	prob.initial_vector, prob.hamiltonian
117	)
118
119	# create the BasisSetRepresentation
120	bsr = BasisSetRepresentation(prob.hamiltonian, addr_or_vec; bsr_kwargs...)	91✔
121	matrix = bsr.sparse_matrix	91✔
122	basis = bsr.basis	91✔
123
124	initial_vector = _set_up_initial_vector(prob.hamiltonian, prob.initial_vector, basis)	237✔
125
126	return IterativeEDSolver(algorithm, prob, matrix, initial_vector, basis, solver_kwargs)	91✔
127	end
128
129	struct DenseEDSolver{P,A,BSR}
130	problem::P	82✔
131	algorithm::A
132	basis_set_rep::BSR
133	solver_kwargs::NamedTuple
134	end
135	function Base.show(io::IO, s::DenseEDSolver)	4✔
136	io = IOContext(io, :compact => true)	4✔
137	print(io, "DenseEDSolver\n for hamiltonian = ")	4✔
138	show(io, s.problem.hamiltonian)	4✔
139	print(io, ",\n kwargs = ")	4✔
140	show(io, s.solver_kwargs)	4✔
141	print(io, "\n)")	4✔
142	end
143
144	function CommonSolve.init(	164✔
145	prob::ExactDiagonalizationProblem, algorithm::LinearAlgebraSolver; kwargs...
146	)
147	# Merge keyword arguments from problem, algorithm and ones passed to this function
148	# and split them into sets that are passed to BasisSetRepresentation and ones that
149	# are left for the solver.
150	kwargs = (; sizelim=1e5, prob.kwargs..., algorithm.kwargs..., kwargs...)	164✔
151	bsr_kwargs, solver_kwargs = split_keys(	82✔
152	kwargs,
153	:sizelim, :cutoff, :filter, :nnzs, :col_hint, :sort, :max_depth, :minimum_size
154	)
155
156	# determine the seed address to build the matrix from
157	addr_or_vec = _set_up_starting_address(	82✔
158	prob.initial_vector, prob.hamiltonian
159	)
160
161	# create the BasisSetRepresentation
162	bsr = BasisSetRepresentation(prob.hamiltonian, addr_or_vec; bsr_kwargs...)	149✔
163
164	return DenseEDSolver(prob, algorithm, bsr, solver_kwargs)	82✔
165	end

RimuQMC / Rimu.jl / 14544069614

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous