12207469360

Committed 06 Dec 2024 11:07PM UTC coverage: 94.667% (+0.2%) from 94.458%

Build # 12207469360

Build Type

push

github

Committed by

web-flow

Commit Message

Fast basis (#287)

# New features

* `build_basis` and `build_sparse_matrix_from_LO` are now parallel.
* `build_basis` and `build_sparse_matrix_from_LO` now support an
argument `max_depth` which limits the number of steps of the breadth
first search to perform.
* New method `build_basis(::AbstractFockAddress)` which is much faster.

# Bugfixes

* `OccupationNumberFS` now has the same ordering as `BoseFS`.
* Fix a bug where `MatrixHamiltonian` would not work with sparse
matrices that have implicit zeros on the diagonal.

Run Details

233 of 241 new or added lines in 5 files covered. (96.68%)

1 existing line in 1 file now uncovered.

6799 of 7182 relevant lines covered (94.67%)

17030919.16 hits per line

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

99.28

/src/ExactDiagonalization/basis_breadth_first_search.jl

const SubVector{T} = SubArray{T,1,Vector{T},Tuple{UnitRange{Int64}},true}
"""
    UniformSplit(array::Vector, min_chunk_size, max_chunks)

Split the array into at most `max_chunks` subarrays, with each at least `min_chunk_size` long.
If the number of elements in `array` is not divisible by the determined chunk size, the leftover elements are placed in the last chunk.

```jldoctest
julia> using Rimu.ExactDiagonalization: UniformSplit

julia> UniformSplit(collect(1:10), 3, 3)
3-element UniformSplit{Int64}:
 [1, 2, 3]
 [4, 5, 6]
 [7, 8, 9, 10]

julia> UniformSplit(collect(1:13), 5, 3)
2-element UniformSplit{Int64}:
 [1, 2, 3, 4, 5, 6]
 [7, 8, 9, 10, 11, 12, 13]
```
"""
struct UniformSplit{T} <: AbstractVector{SubVector{T}}
    array::Vector{T}
    n_chunks::Int
    chunk_size::Int

    function UniformSplit(array::Vector{T}, min_chunk_size, max_chunks) where {T}
        chunk_size = fld(length(array), max_chunks)
        if min_chunk_size > chunk_size
            n_chunks = max(fld(length(array), min_chunk_size), 1)
            chunk_size = fld(length(array), n_chunks)
        else
            n_chunks = max(fld(length(array), chunk_size), 1)
        end
        return new{T}(array, n_chunks, chunk_size)
    end
end

Base.size(split::UniformSplit) = (split.n_chunks,)

function Base.getindex(split::UniformSplit, i)
    index_start = (i - 1) * split.chunk_size + 1
    if i == split.n_chunks
        index_end = length(split.array)
    elseif 0 < i < split.n_chunks
        index_end = i * split.chunk_size
    else
        throw(BoundsError(split, i))
    end
    return view(split.array, index_start:index_end)
end

"""
    bb! = BasisBuilder{A::Type}(; col_hint=0)
    bb!(operator::AbstractOperator, pairs, seen)

Functor used with [`basis_breadth_first_search`](@ref) to build a basis of addresses of type
`A` from an operator. It contains a set of addresses (`bb!.frontier`) that is collected from
the `offdiagonals` of `operator` with all addresses contained the list of address-value
pairs `pairs` that are not element of `seen`.
"""
struct BasisBuilder{A}
    frontier::Set{A}

    function BasisBuilder{A}(; col_hint=0) where {A}
        return new{A}(sizehint!(Set{A}(), col_hint))
    end
end
init_accumulator(::Type{<:BasisBuilder}) = nothing
update_accumulator!(_, ::BasisBuilder, _) = nothing
function finalize_accumulator!(::Nothing, basis, sort; kwargs...)
    if sort
        return sort!(basis; kwargs...)
    else
        return basis
    end
end

function (bb::BasisBuilder)(operator, ks, seen)
    empty!(bb.frontier)
    for (k1, _) in ks
        for (k2, val) in offdiagonals(operator, k1)
            if val ≠ 0 && k2 ∉ seen
                push!(bb.frontier, k2)
            end
        end
    end
end

"""
    mb! = MatrixBuilder{A::Type}(; col_hint=0)
    mb!(operator::AbstractOperator, pairs, seen)

Functor used with [`basis_breadth_first_search`](@ref) to build a matrix from an
operator. It contains a set of addresses (`mb!.frontier`) that is collected from the
`offdiagonals` of `operator` with all addresses contained the list of address-value pairs
`pairs` that are not element of `seen`.

It also collects `mb!.is`, `mb!.js`, and `mb!.vs` which are used to build the sparse matrix
via `sparse!`.
"""
struct MatrixBuilder{A,T,I}
    js::Vector{I}
    is::Vector{A}
    vs::Vector{T}
    column_buffer::Dict{A,T}
    frontier::Set{A}

    function MatrixBuilder{A,T,I}(; col_hint=0) where {A,T,I}
        return new{A,T,I}(
            sizehint!(I[], col_hint),
            sizehint!(A[], col_hint),
            sizehint!(T[], col_hint),
            sizehint!(Dict{A,T}(), col_hint),
            sizehint!(Set{A}(), col_hint),
        )
    end
end

function init_accumulator(::Type{<:MatrixBuilder{<:Any,T,I}}) where {I,T}
    return MatrixBuilderAccumulator{T,I}()
end

function (builder::MatrixBuilder{<:Any,T})(operator, columns, seen) where {T}
    empty!(builder.is)
    empty!(builder.js)
    empty!(builder.vs)
    empty!(builder.frontier)
    for (col_key, col_index) in columns
        empty!(builder.column_buffer)
        diag = diagonal_element(operator, col_key)
        if !iszero(diag)
            builder.column_buffer[col_key] = diag
        end
        for (row_key, value) in offdiagonals(operator, col_key)
            if !iszero(value)
                new_value = get(builder.column_buffer, row_key, zero(T)) + value
                builder.column_buffer[row_key] = new_value
                row_key ∉ seen && push!(builder.frontier, row_key)
            end
        end

        old_len = length(builder.is)
        new_len = length(builder.is) + length(builder.column_buffer)
        resize!(builder.is, new_len)
        resize!(builder.js, new_len)
        resize!(builder.vs, new_len)
        @inbounds for (i, (row_key, value)) in enumerate(builder.column_buffer)
            builder.is[old_len + i] = row_key
            builder.js[old_len + i] = col_index
            builder.vs[old_len + i] = value
        end
    end
end

"""
    struct MatrixBuilderAccumulator

Used in conjunction with [`basis_breadth_first_search`](@ref) and
[`MatrixBuilder`](@ref). It is used to combine the sparse matrix as it is being built by
multiple threads.
"""
struct MatrixBuilderAccumulator{T,I}
    is::Vector{I}
    js::Vector{I}
    vs::Vector{T}

    MatrixBuilderAccumulator{T,I}() where {T,I} = new{T,I}(I[], I[], T[])
end

function update_accumulator!(acc::MatrixBuilderAccumulator, mb::MatrixBuilder, mapping)
    for i in eachindex(mb.is)
        if haskey(mapping, mb.is[i]) # skip keys that were filtered
            push!(acc.is, mapping[mb.is[i]])
            push!(acc.js, mb.js[i])
            push!(acc.vs, mb.vs[i])
        end
    end
end

function finalize_accumulator!(
    acc::MatrixBuilderAccumulator{T,I}, basis, sort; kwargs...
) where {T,I}
    n = length(basis)

    # see docstring of `sparse!` for an explanation for what this is
    klasttouch = Vector{I}(undef, n)
    csrrowptr = Vector{I}(undef, n + 1)
    csrcolval = Vector{I}(undef, length(acc.is))
    csrnzval = Vector{T}(undef, length(acc.is))

    matrix = SparseArrays.sparse!(
        acc.is, acc.js, acc.vs, n, n, +,
        klasttouch, csrrowptr, csrcolval, csrnzval,
        acc.is, acc.js, acc.vs,
    )
    if sort
        # `csrcolval` is no longer needed, so we can reuse it
        perm = resize!(csrcolval, length(basis))
        sortperm!(perm, basis; kwargs...)
        permute!(matrix, perm, perm), permute!(basis, perm)
    else
        return matrix, basis
    end
end

"""
    basis_breadth_first_search(::Type{Builder}, operator, starting_basis)

Internal function that performs breadth-first search (BFS) on an operator.

`Builder` is either [`MatrixBuilder`](@ref) or [`BasisBuilder`](@ref), which triggers
building a matrix and basis, or only the basis of addresses, respectively.
"""
function basis_breadth_first_search(
    ::Type{Builder}, operator, basis::Vector{A};
    min_batch_size=100,
    max_tasks=4 * Threads.nthreads(),

    max_depth=Inf,
    minimum_size=Inf,

    cutoff=nothing,
    filter=isnothing(cutoff) ? Returns(true) : a -> diagonal_element(operator, a) ≤ cutoff,

    sizelim=Inf,
    nnzs=0,
    col_hint=0,

    sort=false,
    by=identity,
    rev=false,
    lt=isless,
    order=Base.Forward,
) where {Builder,A}

    dim = dimension(operator, basis[1])
    if dim > sizelim
        throw(ArgumentError("dimension $dim larger than sizelim $sizelim"))
    end
    sizehint!(basis, nnzs)
    sort_kwargs = (; by, rev, lt, order)

    # addresses already visited and addresses visited in the last layer
    seen = Set(basis)
    last_seen = empty(seen)

    # addresses to visit and their index
    curr_frontier = collect(zip(basis, eachindex(basis)))
    next_frontier = empty(curr_frontier)

    # map from address to index
    mapping = Dict(zip(basis, eachindex(basis)))

    # builders store task-local storage and choose whether a matrix is to be built or not
    builders = [Builder(; col_hint) for _ in 1:max_tasks]
    result_accumulator = init_accumulator(Builder)

    depth = 0
    early_stop = false
    while !isempty(curr_frontier)
        depth += 1
        early_stop = length(basis) > minimum_size || depth > max_depth

        # We can stop here when not constructing the matrix. If we are, we still need
        # to do another round to evaluate the columns.
        if Builder <: BasisBuilder && early_stop
            break
        end

        # Split the workload into chunks and spawn a task for each. These now run
        # asynchronously while the main thread continues.
        split_frontier = UniformSplit(curr_frontier, min_batch_size, max_tasks)
        tasks = map(enumerate(split_frontier)) do (i, sub_frontier)
            Threads.@spawn builders[$i]($operator, $sub_frontier, seen)
        end

        # Procesess the tasks in order they were spawned. This is fine as we've used more
        # tasks than threads.
        for (task, builder) in zip(tasks, builders)
            wait(task)

            result = builder.frontier
            for k in result
                if k ∉ last_seen
                    push!(last_seen, k)
                    if !early_stop && (isnothing(filter) || filter(k))
                        push!(basis, k)
                        push!(next_frontier, (k, length(basis)))
                        mapping[k] = length(basis)
                    end
                end
            end

            # update matrix if needed
            update_accumulator!(result_accumulator, builder, mapping)
        end

        # This had to wait until now so we don't get into data races
        union!(seen, last_seen)
        curr_frontier, next_frontier = next_frontier, curr_frontier
        empty!(next_frontier)
        empty!(last_seen)
    end

    return finalize_accumulator!(result_accumulator, basis, sort; sort_kwargs...)
end

function _address_to_basis(operator, addr_or_basis)
    if addr_or_basis isa AbstractVector || addr_or_basis isa Tuple
        check_address_type(operator, eltype(addr_or_basis))
        return collect(addr_or_basis)
    else
        check_address_type(operator, typeof(addr_or_basis))
        return [addr_or_basis]
    end
end

"""
    build_basis(
        ham, address=starting_address(ham);
        cutoff, filter, sizelim, sort=false, kwargs...
    ) -> basis
    build_basis(ham, addresses::AbstractVector; kwargs...)

Get all basis element of a linear operator `ham` that are reachable (via
non-zero matrix elements) from the address `address`, returned as a vector.
Instead of a single address, a vector of `addresses` can be passed.
Does not return the matrix, for that purpose use [`BasisSetRepresentation`](@ref).

Providing an energy cutoff will skip addresses with diagonal elements greater
than `cutoff`. Alternatively, an arbitrary `filter` function can be used instead.
Addresses passed as arguments are not filtered.

Providing a `max_depth` will limit the size of the basis by only visiting addresses that are
connected to the `starting_address` through `max_depth` hops through the
Hamiltonian. Similarly, providing `minimum_size` will stop the bulding process after the
basis reaches a length of at least `minimum_size`.

A maximum basis size `sizelim` can be set which will throw an error if the expected
dimension of `ham` is larger than `sizelim`. This may be useful when memory may be a
concern. These options are disabled by default.

!!! warning
        The order the basis is returned in is arbitrary and non-deterministic. Use
        `sort=true` if the ordering matters.

"""
function build_basis(operator, addr=starting_address(operator); sizelim=Inf, kwargs...)
    basis = _address_to_basis(operator, addr)
    return basis_breadth_first_search(BasisBuilder{eltype(basis)}, operator, basis; sizelim, kwargs...)
end

"""
    build_sparse_matrix_from_LO(
        ham, address=starting_address(ham);
        cutoff, filter=nothing, nnzs, col_hint, sort=false, kwargs...
    ) -> sparse_matrix, basis
    build_sparse_matrix_from_LO(ham, addresses::AbstractVector; kwargs...)

Create a sparse matrix `sparse_matrix` of all reachable matrix elements of a linear operator
`ham` starting from `address`. Instead of a single address, a vector of `addresses` can be
passed.  The vector `basis` contains the addresses of basis configurations.

Providing the number `nnzs` of expected calculated matrix elements and `col_hint` for the
estimated number of nonzero off-diagonal matrix elements in each matrix column may improve
performance.

Providing an energy cutoff will skip the columns and rows with diagonal elements greater
than `cutoff`. Alternatively, an arbitrary `filter` function can be used instead. These are
not enabled by default. To generate the matrix truncated to the subspace spanned by the
`addresses`, use `filter = Returns(false)`.

Providing a `max_depth` will limit the size of the matrix by only visiting addresses that
are connected to the `starting_address` through `max_depth` hops through the
Hamiltonian. Similarly, providing `minimum_size` will stop the bulding process after the
basis reaches a length of at least `minimum_size`.

Setting `sort` to `true` will sort the `basis` and order the matrix rows and columns
accordingly. This is useful when the order of the columns matters, e.g. when comparing
matrices. Any additional keyword arguments are passed on to `Base.sortperm`.

!!! warning
        The order of the returned rows and columns is arbitrary and non-deterministic. Use
        `sort=true` if the ordering matters.

See [`BasisSetRepresentation`](@ref).
"""
function build_sparse_matrix_from_LO(
    operator, addr=starting_address(operator); sizelim=1e7, kwargs...
)
    basis = _address_to_basis(operator, addr)
    T = eltype(operator)
    return basis_breadth_first_search(
        MatrixBuilder{eltype(basis),T,Int32}, operator, basis;
        sizelim, kwargs...
    )
end

1	const SubVector{T} = SubArray{T,1,Vector{T},Tuple{UnitRange{Int64}},true}
2	"""
3	UniformSplit(array::Vector, min_chunk_size, max_chunks)
4
5	Split the array into at most `max_chunks` subarrays, with each at least `min_chunk_size` long.
6	If the number of elements in `array` is not divisible by the determined chunk size, the leftover elements are placed in the last chunk.
7
8	```jldoctest
9	julia> using Rimu.ExactDiagonalization: UniformSplit
10
11	julia> UniformSplit(collect(1:10), 3, 3)
12	3-element UniformSplit{Int64}:
13	[1, 2, 3]
14	[4, 5, 6]
15	[7, 8, 9, 10]
16
17	julia> UniformSplit(collect(1:13), 5, 3)
18	2-element UniformSplit{Int64}:
19	[1, 2, 3, 4, 5, 6]
20	[7, 8, 9, 10, 11, 12, 13]
21	```
22	"""
23	struct UniformSplit{T} <: AbstractVector{SubVector{T}}
24	array::Vector{T}
25	n_chunks::Int
26	chunk_size::Int
27
28	function UniformSplit(array::Vector{T}, min_chunk_size, max_chunks) where {T}	2,266✔
29	chunk_size = fld(length(array), max_chunks)	2,266✔
30	if min_chunk_size > chunk_size	2,266✔
31	n_chunks = max(fld(length(array), min_chunk_size), 1)	2,250✔
32	chunk_size = fld(length(array), n_chunks)	2,250✔
33	else
34	n_chunks = max(fld(length(array), chunk_size), 1)	16✔
35	end
36	return new{T}(array, n_chunks, chunk_size)	2,266✔
37	end
38	end
39
40	Base.size(split::UniformSplit) = (split.n_chunks,)	4,600✔
41
42	function Base.getindex(split::UniformSplit, i)	2,442✔
43	index_start = (i - 1) * split.chunk_size + 1	2,442✔
44	if i == split.n_chunks	2,442✔
45	index_end = length(split.array)	2,278✔
46	elseif 0 < i < split.n_chunks	164✔
47	index_end = i * split.chunk_size	164✔
48	else
NEW 49	throw(BoundsError(split, i))	×
50	end
51	return view(split.array, index_start:index_end)	2,442✔
52	end
53
54	"""
55	bb! = BasisBuilder{A::Type}(; col_hint=0)
56	bb!(operator::AbstractOperator, pairs, seen)
57
58	Functor used with [`basis_breadth_first_search`](@ref) to build a basis of addresses of type
59	`A` from an operator. It contains a set of addresses (`bb!.frontier`) that is collected from
60	the `offdiagonals` of `operator` with all addresses contained the list of address-value
61	pairs `pairs` that are not element of `seen`.
62	"""
63	struct BasisBuilder{A}
64	frontier::Set{A}
65
66	function BasisBuilder{A}(; col_hint=0) where {A}	8,208✔
67	return new{A}(sizehint!(Set{A}(), col_hint))	5,472✔
68	end
69	end
70	init_accumulator(::Type{<:BasisBuilder}) = nothing	342✔
71	update_accumulator!(_, ::BasisBuilder, _) = nothing	878✔
72	function finalize_accumulator!(::Nothing, basis, sort; kwargs...)	696✔
73	if sort	342✔
74	return sort!(basis; kwargs...)	32✔
75	else
76	return basis	322✔
77	end
78	end
79
80	function (bb::BasisBuilder)(operator, ks, seen)	877✔
81	empty!(bb.frontier)	562,188✔
82	for (k1, _) in ks	1,756✔
83	for (k2, val) in offdiagonals(operator, k1)	50,702✔
84	if val ≠ 0 && k2 ∉ seen	400,704✔
85	push!(bb.frontier, k2)	102,685✔
86	end
87	end	383,624✔
88	end	25,754✔
89	end
90
91	"""
92	mb! = MatrixBuilder{A::Type}(; col_hint=0)
93	mb!(operator::AbstractOperator, pairs, seen)
94
95	Functor used with [`basis_breadth_first_search`](@ref) to build a matrix from an
96	operator. It contains a set of addresses (`mb!.frontier`) that is collected from the
97	`offdiagonals` of `operator` with all addresses contained the list of address-value pairs
98	`pairs` that are not element of `seen`.
99
100	It also collects `mb!.is`, `mb!.js`, and `mb!.vs` which are used to build the sparse matrix
101	via `sparse!`.
102	"""
103	struct MatrixBuilder{A,T,I}
104	js::Vector{I}
105	is::Vector{A}
106	vs::Vector{T}
107	column_buffer::Dict{A,T}
108	frontier::Set{A}
109
110	function MatrixBuilder{A,T,I}(; col_hint=0) where {A,T,I}	5,440✔
111	return new{A,T,I}(	2,720✔
112	sizehint!(I[], col_hint),
113	sizehint!(A[], col_hint),
114	sizehint!(T[], col_hint),
115	sizehint!(Dict{A,T}(), col_hint),
116	sizehint!(Set{A}(), col_hint),
117	)
118	end
119	end
120
121	function init_accumulator(::Type{<:MatrixBuilder{<:Any,T,I}}) where {I,T}	340✔
122	return MatrixBuilderAccumulator{T,I}()	340✔
123	end
124
125	function (builder::MatrixBuilder{<:Any,T})(operator, columns, seen) where {T}	1,524✔
126	empty!(builder.is)	125,164✔
127	empty!(builder.js)	125,207✔
128	empty!(builder.vs)	125,202✔
129	empty!(builder.frontier)	79,902✔
130	for (col_key, col_index) in columns	3,048✔
131	empty!(builder.column_buffer)	987,106✔
132	diag = diagonal_element(operator, col_key)	14,715✔
133	if !iszero(diag)	13,866✔
134	builder.column_buffer[col_key] = diag	12,092✔
135	end
136	for (row_key, value) in offdiagonals(operator, col_key)	25,710✔
137	if !iszero(value)	2,158,468✔
138	new_value = get(builder.column_buffer, row_key, zero(T)) + value	605,936✔
139	builder.column_buffer[row_key] = new_value	594,029✔
140	row_key ∉ seen && push!(builder.frontier, row_key)	1,136,509✔
141	end
142	end	4,306,153✔
143
144	old_len = length(builder.is)	13,849✔
145	new_len = length(builder.is) + length(builder.column_buffer)	13,849✔
146	resize!(builder.is, new_len)	13,849✔
147	resize!(builder.js, new_len)	13,858✔
148	resize!(builder.vs, new_len)	13,861✔
149	@inbounds for (i, (row_key, value)) in enumerate(builder.column_buffer)	27,722✔
150	builder.is[old_len + i] = row_key	235,971✔
151	builder.js[old_len + i] = col_index	235,985✔
152	builder.vs[old_len + i] = value	236,019✔
153	end	458,135✔
154	end	13,861✔
155	end
156
157	"""
158	struct MatrixBuilderAccumulator
159
160	Used in conjunction with [`basis_breadth_first_search`](@ref) and
161	[`MatrixBuilder`](@ref). It is used to combine the sparse matrix as it is being built by
162	multiple threads.
163	"""
164	struct MatrixBuilderAccumulator{T,I}
165	is::Vector{I}
166	js::Vector{I}
167	vs::Vector{T}
168
169	MatrixBuilderAccumulator{T,I}() where {T,I} = new{T,I}(I[], I[], T[])	340✔
170	end
171
172	function update_accumulator!(acc::MatrixBuilderAccumulator, mb::MatrixBuilder, mapping)	1,524✔
173	for i in eachindex(mb.is)	1,524✔
174	if haskey(mapping, mb.is[i]) # skip keys that were filtered	423,286✔
175	push!(acc.is, mapping[mb.is[i]])	236,606✔
176	push!(acc.js, mb.js[i])	236,606✔
177	push!(acc.vs, mb.vs[i])	236,606✔
178	end
179	end	236,982✔
180	end
181
182	function finalize_accumulator!(	680✔
183	acc::MatrixBuilderAccumulator{T,I}, basis, sort; kwargs...
184	) where {T,I}
185	n = length(basis)	340✔
186
187	# see docstring of `sparse!` for an explanation for what this is
188	klasttouch = Vector{I}(undef, n)	680✔
189	csrrowptr = Vector{I}(undef, n + 1)	680✔
190	csrcolval = Vector{I}(undef, length(acc.is))	678✔
191	csrnzval = Vector{T}(undef, length(acc.is))	678✔
192
193	matrix = SparseArrays.sparse!(	340✔
194	acc.is, acc.js, acc.vs, n, n, +,
195	klasttouch, csrrowptr, csrcolval, csrnzval,
196	acc.is, acc.js, acc.vs,
197	)
198	if sort	340✔
199	# `csrcolval` is no longer needed, so we can reuse it
200	perm = resize!(csrcolval, length(basis))	54✔
201	sortperm!(perm, basis; kwargs...)	68✔
202	permute!(matrix, perm, perm), permute!(basis, perm)	54✔
203	else
204	return matrix, basis	286✔
205	end
206	end
207
208	"""
209	basis_breadth_first_search(::Type{Builder}, operator, starting_basis)
210
211	Internal function that performs breadth-first search (BFS) on an operator.
212
213	`Builder` is either [`MatrixBuilder`](@ref) or [`BasisBuilder`](@ref), which triggers
214	building a matrix and basis, or only the basis of addresses, respectively.
215	"""
216	function basis_breadth_first_search(	1,372✔
217	::Type{Builder}, operator, basis::Vector{A};
218	min_batch_size=100,
219	max_tasks=4 * Threads.nthreads(),
220
221	max_depth=Inf,
222	minimum_size=Inf,
223
224	cutoff=nothing,
225	filter=isnothing(cutoff) ? Returns(true) : a -> diagonal_element(operator, a) ≤ cutoff,	516✔
226
227	sizelim=Inf,
228	nnzs=0,
229	col_hint=0,
230
231	sort=false,
232	by=identity,
233	rev=false,
234	lt=isless,
235	order=Base.Forward,
236	) where {Builder,A}
237
238	dim = dimension(operator, basis[1])	686✔
239	if dim > sizelim	1,292✔
240	throw(ArgumentError("dimension $dim larger than sizelim $sizelim"))	4✔
241	end
242	sizehint!(basis, nnzs)	682✔
243	sort_kwargs = (; by, rev, lt, order)	682✔
244
245	# addresses already visited and addresses visited in the last layer
246	seen = Set(basis)	682✔
247	last_seen = empty(seen)	682✔
248
249	# addresses to visit and their index
250	curr_frontier = collect(zip(basis, eachindex(basis)))	682✔
251	next_frontier = empty(curr_frontier)	682✔
252
253	# map from address to index
254	mapping = Dict(zip(basis, eachindex(basis)))	682✔
255
256	# builders store task-local storage and choose whether a matrix is to be built or not
257	builders = [Builder(; col_hint) for _ in 1:max_tasks]	682✔
258	result_accumulator = init_accumulator(Builder)	682✔
259
260	depth = 0	682✔
261	early_stop = false	682✔
262	while !isempty(curr_frontier)	2,944✔
263	depth += 1	2,286✔
264	early_stop = length(basis) > minimum_size \|\| depth > max_depth	4,558✔
265
266	# We can stop here when not constructing the matrix. If we are, we still need
267	# to do another round to evaluate the columns.
268	if Builder <: BasisBuilder && early_stop	2,286✔
269	break	24✔
270	end
271
272	# Split the workload into chunks and spawn a task for each. These now run
273	# asynchronously while the main thread continues.
274	split_frontier = UniformSplit(curr_frontier, min_batch_size, max_tasks)	2,262✔
275	tasks = map(enumerate(split_frontier)) do (i, sub_frontier)	2,262✔
276	Threads.@spawn builders[$i]($operator, $sub_frontier, seen)	4,802✔
277	end
278
279	# Procesess the tasks in order they were spawned. This is fine as we've used more
280	# tasks than threads.
281	for (task, builder) in zip(tasks, builders)	4,524✔
282	wait(task)	2,402✔
283
284	result = builder.frontier	2,402✔
285	for k in result	4,804✔
286	if k ∉ last_seen	171,122✔
287	push!(last_seen, k)	39,454✔
288	if !early_stop && (isnothing(filter) \|\| filter(k))	39,466✔
289	push!(basis, k)	39,200✔
290	push!(next_frontier, (k, length(basis)))	39,200✔
291	mapping[k] = length(basis)	39,200✔
292	end
293	end
294	end	179,652✔
295
296	# update matrix if needed
297	update_accumulator!(result_accumulator, builder, mapping)	2,402✔
298	end	2,542✔
299
300	# This had to wait until now so we don't get into data races
301	union!(seen, last_seen)	2,262✔
302	curr_frontier, next_frontier = next_frontier, curr_frontier	2,262✔
303	empty!(next_frontier)	40,138✔
304	empty!(last_seen)	249,960✔
305	end	2,262✔
306
307	return finalize_accumulator!(result_accumulator, basis, sort; sort_kwargs...)	682✔
308	end
309
310	function _address_to_basis(operator, addr_or_basis)	688✔
311	if addr_or_basis isa AbstractVector \|\| addr_or_basis isa Tuple	688✔
312	check_address_type(operator, eltype(addr_or_basis))	16✔
313	return collect(addr_or_basis)	32✔
314	else
315	check_address_type(operator, typeof(addr_or_basis))	688✔
316	return [addr_or_basis]	670✔
317	end
318	end
319
320	"""
321	build_basis(
322	ham, address=starting_address(ham);
323	cutoff, filter, sizelim, sort=false, kwargs...
324	) -> basis
325	build_basis(ham, addresses::AbstractVector; kwargs...)
326
327	Get all basis element of a linear operator `ham` that are reachable (via
328	non-zero matrix elements) from the address `address`, returned as a vector.
329	Instead of a single address, a vector of `addresses` can be passed.
330	Does not return the matrix, for that purpose use [`BasisSetRepresentation`](@ref).
331
332	Providing an energy cutoff will skip addresses with diagonal elements greater
333	than `cutoff`. Alternatively, an arbitrary `filter` function can be used instead.
334	Addresses passed as arguments are not filtered.
335
336	Providing a `max_depth` will limit the size of the basis by only visiting addresses that are
337	connected to the `starting_address` through `max_depth` hops through the
338	Hamiltonian. Similarly, providing `minimum_size` will stop the bulding process after the
339	basis reaches a length of at least `minimum_size`.
340
341	A maximum basis size `sizelim` can be set which will throw an error if the expected
342	dimension of `ham` is larger than `sizelim`. This may be useful when memory may be a
343	concern. These options are disabled by default.
344
345	!!! warning
346	The order the basis is returned in is arbitrary and non-deterministic. Use
347	`sort=true` if the ordering matters.
348
349	"""
350	function build_basis(operator, addr=starting_address(operator); sizelim=Inf, kwargs...)	732✔
351	basis = _address_to_basis(operator, addr)	348✔
352	return basis_breadth_first_search(BasisBuilder{eltype(basis)}, operator, basis; sizelim, kwargs...)	344✔
353	end
354
355	"""
356	build_sparse_matrix_from_LO(
357	ham, address=starting_address(ham);
358	cutoff, filter=nothing, nnzs, col_hint, sort=false, kwargs...
359	) -> sparse_matrix, basis
360	build_sparse_matrix_from_LO(ham, addresses::AbstractVector; kwargs...)
361
362	Create a sparse matrix `sparse_matrix` of all reachable matrix elements of a linear operator
363	`ham` starting from `address`. Instead of a single address, a vector of `addresses` can be
364	passed. The vector `basis` contains the addresses of basis configurations.
365
366	Providing the number `nnzs` of expected calculated matrix elements and `col_hint` for the
367	estimated number of nonzero off-diagonal matrix elements in each matrix column may improve
368	performance.
369
370	Providing an energy cutoff will skip the columns and rows with diagonal elements greater
371	than `cutoff`. Alternatively, an arbitrary `filter` function can be used instead. These are
372	not enabled by default. To generate the matrix truncated to the subspace spanned by the
373	`addresses`, use `filter = Returns(false)`.
374
375	Providing a `max_depth` will limit the size of the matrix by only visiting addresses that
376	are connected to the `starting_address` through `max_depth` hops through the
377	Hamiltonian. Similarly, providing `minimum_size` will stop the bulding process after the
378	basis reaches a length of at least `minimum_size`.
379
380	Setting `sort` to `true` will sort the `basis` and order the matrix rows and columns
381	accordingly. This is useful when the order of the columns matters, e.g. when comparing
382	matrices. Any additional keyword arguments are passed on to `Base.sortperm`.
383
384	!!! warning
385	The order of the returned rows and columns is arbitrary and non-deterministic. Use
386	`sort=true` if the ordering matters.
387
388	See [`BasisSetRepresentation`](@ref).
389	"""
390	function build_sparse_matrix_from_LO(	716✔
391	operator, addr=starting_address(operator); sizelim=1e7, kwargs...
392	)
393	basis = _address_to_basis(operator, addr)	372✔
394	T = eltype(operator)	342✔
395	return basis_breadth_first_search(	342✔
396	MatrixBuilder{eltype(basis),T,Int32}, operator, basis;
397	sizelim, kwargs...
398	)
399	end

RimuQMC / Rimu.jl / 12207469360

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