12002034438

Committed 25 Nov 2024 02:22AM UTC coverage: 94.606% (+0.1%) from 94.481%

Build # 12002034438

Build Type

Pull #287

github

Committed by

mtsch

Commit Message

remove space from doctest

Pull Request Pull Request #287: Fast basis

Run Details

224 of 237 new or added lines in 4 files covered. (94.51%)

3 existing lines in 1 file now uncovered.

6735 of 7119 relevant lines covered (94.61%)

16618298.61 hits per line

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

86.52

/src/ExactDiagonalization/basis_fock.jl

"""
    build_basis(addr::AbstractFockAddress)
    build_basis(::Type{<:AbstractFockAddress}) -> basis

Return all possible Fock states of a given type as a vector. This method is _much_ faster
than `build_basis(::AbstractHamiltonian, ...)`, but does not take matrix blocking into
account. This version of `build_basis` accepts no additional arguments.

All address types except [`OccupationNumberFS`](@ref Main.Rimu.OccupationNumberFS) are
supported.

Returns a sorted vector of length equal to the [`dimension`](@ref) of `addr`.

See also [`AbstractFockAddress`](@ref).
"""
function build_basis(addr::AbstractFockAddress)
    return build_basis(typeof(addr))
end

###
### BoseFS
###
# Algorithm:
# Given an address type of N particles in M modes,
# For all n_1 ∈ (0 ... N)
# - put n_1 particles in the first mode
# - recursively call the function to put the remaining (N-n_1) particles into (M-1) modes.
# The recursion is parallelised through spawning tasks.

# this is equivalent to `dimension(BoseFS{n,m})`, but does not return a `BigInt`.
_bose_dimension(n, m) = binomial(n + m - 1, n)

function build_basis(::Type{<:BoseFS{N,M}}) where {N,M}
    T = typeof(near_uniform(BoseFS{N,M}))
    result = Vector{T}(undef, _bose_dimension(N, M))
    _bose_basis!(result, (), 1, N, Val(M), 2 * Threads.nthreads())
    return result
end

# result: the basis that is eventually returned
# postfix: partially build ONR
# index: the position the built address is written to in result
# remaining_n: the number of particles to be placed in the remaining part of the ONR
# M: the number of modes in the ONR left to fill
# n_tasks: number of tasks to spawn to build the basis in parallel
@inline function _bose_basis!(
    result::Vector, postfix, index, remaining_n, ::Val{M}, n_tasks::Int
) where {M}
    @sync if M < 5 || remaining_n ≤ 1 || n_tasks ≤ 0
        _bose_basis!(result, postfix, index, remaining_n, Val(M))
    else
        n_tasks ÷= 2
        Threads.@spawn begin
            _bose_basis!(result, $(0, postfix...), $index, $remaining_n, Val(M-1), $n_tasks)
        end
        index += _bose_dimension(remaining_n, M - 1)
        Threads.@spawn begin
            _bose_basis!(
                result, $(1, postfix...), $index, $(remaining_n - 1), Val(M-1), $n_tasks
            )
        end
        index += _bose_dimension(remaining_n - 1, M - 1)
        for n in 2:remaining_n
            _bose_basis!(result, (n, postfix...), index, remaining_n - n, Val(M-1))
            index += _bose_dimension(remaining_n - n, M - 1)
        end
    end
end

@inline function _bose_basis!(
    result::Vector{T}, postfix, index, remaining_n, ::Val{M}
) where {M,T}
    if remaining_n == 0
        @inbounds result[index] = T((ntuple(Returns(0), Val(M))..., postfix...))
    elseif remaining_n == 1
        _basis_basecase_N1!(result, postfix, index, Val(M))
    elseif M == 1
        @inbounds result[index] = T((remaining_n, postfix...))
    elseif M == 2
        _bose_basis_basecase_M2!(result, postfix, index, remaining_n)
    elseif M == 3
        _bose_basis_basecase_M3!(result, postfix, index, remaining_n)
    else
        for n in 0:remaining_n
            _bose_basis!(result, (n, postfix...), index, remaining_n - n, Val(M - 1))
            index += _bose_dimension(remaining_n - n, M - 1)
        end
    end
    return nothing
end

###
### FermiFS
###
# The algorithm is similar to the BoseFS one.

# this is equivalent to `dimension(FermiFS{n,m})`, but does not return a `BigInt`.
_fermi_dimension(n, m) = binomial(m, n)

function build_basis(::Type{<:FermiFS{N,M}}) where {N,M}
    T = typeof(near_uniform(FermiFS{N,M}))
    result = Vector{T}(undef, _fermi_dimension(N, M))
    _fermi_basis!(result, (), 1, N, Val(M), 2 * Threads.nthreads())
    return result
end

# result: the basis that is eventually returned
# postfix: partially build ONR
# index: the position the built address is written to in result
# remaining_n: the number of particles to be placed in the remaining part of the ONR
# M: the number of modes in the ONR left to fill
# n_tasks: number of tasks to spawn to build the basis in parallel
@inline function _fermi_basis!(
    result::Vector, postfix, index, remaining_n, ::Val{M}, n_tasks
) where {M}
    @sync if M < 5 || remaining_n ≤ M || remaining_n == 1 || n_tasks ≤ 0
        _fermi_basis!(result, postfix, index, remaining_n, Val(M))
    else
        n_tasks ÷= 2
        Threads.@spawn begin
            _fermi_basis!(result, $(0, postfix...), $index, $remaining_n, Val(M-1), $n_tasks)
        end
        index += _fermi_dimension(remaining_n, M - 1)
        _fermi_basis!(result, (1, postfix...), index, remaining_n - 1, Val(M - 1), n_tasks)
    end
end

@inline function _fermi_basis!(
    result::Vector{T}, postfix, index, remaining_n, ::Val{M}
) where {M,T}
    @assert M ≥ remaining_n
    if remaining_n == 0
        @inbounds result[index] = T((ntuple(Returns(0), Val(M))..., postfix...))
    elseif remaining_n == M
        @inbounds result[index] = T((ntuple(Returns(1), Val(M))..., postfix...))
    elseif remaining_n == 1
        _basis_basecase_N1!(result, postfix, index, Val(M))
    else
        _fermi_basis!(result, (0, postfix...), index, remaining_n, Val(M - 1))
        index += _fermi_dimension(remaining_n, M - 1)
        _fermi_basis!(result, (1, postfix...), index, remaining_n - 1, Val(M - 1))
    end
    return nothing
end

###
### CompositeFS
###
function build_basis(::Type{C}) where {T,C<:CompositeFS{<:Any,<:Any,<:Any,T}}
    sub_results = map(build_basis, reverse(Tuple(T.parameters)))
    result = Vector{C}(undef, prod(length, sub_results))
    Threads.@threads for i in eachindex(result)
        @inbounds result[i] = C(_collect_addrs(sub_results, i))
    end
    return result
end

@inline _collect_addrs(::Tuple{}, _) = ()
@inline function _collect_addrs((v, vs...), i)
    rest, curr = fldmod1(i, length(v))
    return (_collect_addrs(vs, rest)..., v[curr])
end

###
### Base cases
###
# Base case for 1 particle in M modes.
@inline function _basis_basecase_N1!(
    result::Vector{T}, postfix, index, ::Val{M}
) where {M,T}
    rest = ntuple(Returns(0), Val(M))
    for k in 1:M
        @inbounds result[index + k - 1] = T((setindex(rest, 1, k)..., postfix...))
    end
end
# Base case for bosons with `remaining_n` particles in 2 modes.
@inline function _bose_basis_basecase_M2!(
    result::Vector{T}, postfix, index, remaining_n
) where {T}
    for n1 in 0:remaining_n
        @inbounds result[index + n1] = T((remaining_n - n1, n1, postfix...))
    end
end
# Base case for bosons with `remaining_n` particles in 3 modes.
@inline function _bose_basis_basecase_M3!(
    result::Vector{T}, postfix, index, remaining_n
) where {T}
    k = 0
    for n1 in 0:remaining_n, n2 in 0:(remaining_n - n1)
        @inbounds result[index + k] = T((remaining_n - n1 - n2, n2, n1, postfix...))
        k += 1
    end
end

1	"""
2	build_basis(addr::AbstractFockAddress)
3	build_basis(::Type{<:AbstractFockAddress}) -> basis
4
5	Return all possible Fock states of a given type as a vector. This method is _much_ faster
6	than `build_basis(::AbstractHamiltonian, ...)`, but does not take matrix blocking into
7	account. This version of `build_basis` accepts no additional arguments.
8
9	All address types except [`OccupationNumberFS`](@ref Main.Rimu.OccupationNumberFS) are
10	supported.
11
12	Returns a sorted vector of length equal to the [`dimension`](@ref) of `addr`.
13
14	See also [`AbstractFockAddress`](@ref).
15	"""
16	function build_basis(addr::AbstractFockAddress)	10✔
17	return build_basis(typeof(addr))	10✔
18	end
19
20	###
21	### BoseFS
22	###
23	# Algorithm:
24	# Given an address type of N particles in M modes,
25	# For all n_1 ∈ (0 ... N)
26	# - put n_1 particles in the first mode
27	# - recursively call the function to put the remaining (N-n_1) particles into (M-1) modes.
28	# The recursion is parallelised through spawning tasks.
29
30	# this is equivalent to `dimension(BoseFS{n,m})`, but does not return a `BigInt`.
31	_bose_dimension(n, m) = binomial(n + m - 1, n)	4,562✔
32
33	function build_basis(::Type{<:BoseFS{N,M}}) where {N,M}	10✔
34	T = typeof(near_uniform(BoseFS{N,M}))	10✔
35	result = Vector{T}(undef, _bose_dimension(N, M))	10✔
36	_bose_basis!(result, (), 1, N, Val(M), 2 * Threads.nthreads())	30✔
37	return result	10✔
38	end
39
40	# result: the basis that is eventually returned
41	# postfix: partially build ONR
42	# index: the position the built address is written to in result
43	# remaining_n: the number of particles to be placed in the remaining part of the ONR
44	# M: the number of modes in the ONR left to fill
45	# n_tasks: number of tasks to spawn to build the basis in parallel
46	@inline function _bose_basis!(	38✔
47	result::Vector, postfix, index, remaining_n, ::Val{M}, n_tasks::Int
48	) where {M}
49	@sync if M < 5 \|\| remaining_n ≤ 1 \|\| n_tasks ≤ 0	58✔
50	_bose_basis!(result, postfix, index, remaining_n, Val(M))	44✔
51	else
52	n_tasks ÷= 2	14✔
53	Threads.@spawn begin	28✔
54	_bose_basis!(result, $(0, postfix...), $index, $remaining_n, Val(M-1), $n_tasks)	14✔
55	end
56	index += _bose_dimension(remaining_n, M - 1)	14✔
57	Threads.@spawn begin	28✔
58	_bose_basis!(	14✔
59	result, $(1, postfix...), $index, $(remaining_n - 1), Val(M-1), $n_tasks
60	)
61	end
62	index += _bose_dimension(remaining_n - 1, M - 1)	14✔
63	for n in 2:remaining_n	14✔
64	_bose_basis!(result, (n, postfix...), index, remaining_n - n, Val(M-1))	102✔
65	index += _bose_dimension(remaining_n - n, M - 1)	102✔
66	end	190✔
67	end
68	end
69
70	@inline function _bose_basis!(	4,548✔
71	result::Vector{T}, postfix, index, remaining_n, ::Val{M}
72	) where {M,T}
73	if remaining_n == 0	4,548✔
74	@inbounds result[index] = T((ntuple(Returns(0), Val(M))..., postfix...))	990✔
75	elseif remaining_n == 1	3,558✔
76	_basis_basecase_N1!(result, postfix, index, Val(M))	1,006✔
77	elseif M == 1	2,564✔
NEW 78	@inbounds result[index] = T((remaining_n, postfix...))	×
79	elseif M == 2	2,564✔
NEW 80	_bose_basis_basecase_M2!(result, postfix, index, remaining_n)	×
81	elseif M == 3	2,564✔
82	_bose_basis_basecase_M3!(result, postfix, index, remaining_n)	1,596✔
83	else
84	for n in 0:remaining_n	976✔
85	_bose_basis!(result, (n, postfix...), index, remaining_n - n, Val(M - 1))	4,422✔
86	index += _bose_dimension(remaining_n - n, M - 1)	4,422✔
87	end	4,422✔
88	end
89	return nothing	4,548✔
90	end
91
92	###
93	### FermiFS
94	###
95	# The algorithm is similar to the BoseFS one.
96
97	# this is equivalent to `dimension(FermiFS{n,m})`, but does not return a `BigInt`.
98	_fermi_dimension(n, m) = binomial(m, n)	54✔
99
100	function build_basis(::Type{<:FermiFS{N,M}}) where {N,M}	8✔
101	T = typeof(near_uniform(FermiFS{N,M}))	8✔
102	result = Vector{T}(undef, _fermi_dimension(N, M))	8✔
103	_fermi_basis!(result, (), 1, N, Val(M), 2 * Threads.nthreads())	8✔
104	return result	8✔
105	end
106
107	# result: the basis that is eventually returned
108	# postfix: partially build ONR
109	# index: the position the built address is written to in result
110	# remaining_n: the number of particles to be placed in the remaining part of the ONR
111	# M: the number of modes in the ONR left to fill
112	# n_tasks: number of tasks to spawn to build the basis in parallel
113	@inline function _fermi_basis!(	8✔
114	result::Vector, postfix, index, remaining_n, ::Val{M}, n_tasks
115	) where {M}
116	@sync if M < 5 \|\| remaining_n ≤ M \|\| remaining_n == 1 \|\| n_tasks ≤ 0	8✔
117	_fermi_basis!(result, postfix, index, remaining_n, Val(M))	8✔
118	else
NEW 119	n_tasks ÷= 2	×
NEW 120	Threads.@spawn begin	×
121	_fermi_basis!(result, $(0, postfix...), $index, $remaining_n, Val(M-1), $n_tasks)
122	end
NEW 123	index += _fermi_dimension(remaining_n, M - 1)	×
NEW 124	_fermi_basis!(result, (1, postfix...), index, remaining_n - 1, Val(M - 1), n_tasks)	×
125	end
126	end
127
128	@inline function _fermi_basis!(	100✔
129	result::Vector{T}, postfix, index, remaining_n, ::Val{M}
130	) where {M,T}
131	@assert M ≥ remaining_n	100✔
132	if remaining_n == 0	100✔
NEW 133	@inbounds result[index] = T((ntuple(Returns(0), Val(M))..., postfix...))	×
134	elseif remaining_n == M	100✔
135	@inbounds result[index] = T((ntuple(Returns(1), Val(M))..., postfix...))	20✔
136	elseif remaining_n == 1	80✔
137	_basis_basecase_N1!(result, postfix, index, Val(M))	34✔
138	else
139	_fermi_basis!(result, (0, postfix...), index, remaining_n, Val(M - 1))	46✔
140	index += _fermi_dimension(remaining_n, M - 1)	46✔
141	_fermi_basis!(result, (1, postfix...), index, remaining_n - 1, Val(M - 1))	46✔
142	end
143	return nothing	100✔
144	end
145
146	###
147	### CompositeFS
148	###
149	function build_basis(::Type{C}) where {T,C<:CompositeFS{<:Any,<:Any,<:Any,T}}	4✔
150	sub_results = map(build_basis, reverse(Tuple(T.parameters)))	4✔
151	result = Vector{C}(undef, prod(length, sub_results))	8✔
152	Threads.@threads for i in eachindex(result)	4✔
153	@inbounds result[i] = C(_collect_addrs(sub_results, i))	1,248✔
NEW 154	end	×
155	return result	4✔
156	end
157
158	@inline _collect_addrs(::Tuple{}, _) = ()	1,248✔
159	@inline function _collect_addrs((v, vs...), i)	3,792✔
160	rest, curr = fldmod1(i, length(v))	3,792✔
161	return (_collect_addrs(vs, rest)..., v[curr])	3,792✔
162	end
163
164	###
165	### Base cases
166	###
167	# Base case for 1 particle in M modes.
168	@inline function _basis_basecase_N1!(	1,028✔
169	result::Vector{T}, postfix, index, ::Val{M}
170	) where {M,T}
171	rest = ntuple(Returns(0), Val(M))	1,028✔
172	for k in 1:M	1,028✔
173	@inbounds result[index + k - 1] = T((setindex(rest, 1, k)..., postfix...))	3,522✔
174	end	3,534✔
175	end
176	# Base case for bosons with `remaining_n` particles in 2 modes.
NEW 177	@inline function _bose_basis_basecase_M2!(	×
178	result::Vector{T}, postfix, index, remaining_n
179	) where {T}
NEW 180	for n1 in 0:remaining_n	×
NEW 181	@inbounds result[index + n1] = T((remaining_n - n1, n1, postfix...))	×
NEW 182	end	×
183	end
184	# Base case for bosons with `remaining_n` particles in 3 modes.
185	@inline function _bose_basis_basecase_M3!(	1,588✔
186	result::Vector{T}, postfix, index, remaining_n
187	) where {T}
188	k = 0	1,588✔
189	for n1 in 0:remaining_n, n2 in 0:(remaining_n - n1)	1,588✔
190	@inbounds result[index + k] = T((remaining_n - n1 - n2, n2, n1, postfix...))	18,504✔
191	k += 1	18,504✔
192	end	18,512✔
193	end

joachimbrand / Rimu.jl / 12002034438

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