Hi, I am writing some code to update the columns in an array by repeatedly solving a linear system. I am trying to allocate arrays for the solutions, left-hand-side, and right-hand-side at the beginning of the code and then mutate them to increase performance. Below is a minimal working example (for the use case n will be much larger)
using LinearAlgebra, Distributions, Random, LinearSolve, BenchmarkTools, InvertedIndices
mutable struct G_hat{M<:Matrix{Float64}} # creating a mutable struct with a guess matrix we will update
mat::M
end
function G_hat(N::Int64) # outer constructor function for our guess
G_hat(ones(N,N))
end
struct B_struct{A<:Matrix{Float64}} # immutable struct object for S
mat::A
end
function make_B(N::Int64) # outer constructor function for B_mat
B_struct(rand(N,N))
end
function init_lhs(n::Int64) # outer constructor function for lhs
Vector{Float64}(undef, n-1)
end
f(a,b) = a .- b
function update_lhs!(M, B_j_sl, G_minus_j, B_j)
A = similar(B_j_sl)
mul!(A, G_minus_j', B_j)
M .= B_j_sl - A
return nothing
end
function init_rhs(n::Int64) # outer constructor function for inner prod starting in place one!
Matrix{Float64}(undef, (n-1, n-1))
end
function update_rhs!(M, K, B_sl)
M .= f.(K, B_sl)
return nothing
end
make_ols_estimates! = function(b_hat, G_minus_j, B_minus_j, B_j, K, j, lhs, rhs)
update_lhs!(lhs, @view(B_j[Not(j),:]), G_minus_j, B_j) # update lhs - lots of allocations for some reason
update_rhs!(rhs, K, @view(B_minus_j[Not(j),:])) # update rhs - still lots of allocations!
b_hat .= solve(LinearProblem(rhs, lhs)).u
return nothing
end
sweep_column_two! = function(G_hat_1, B, K, j,p, b_hat, lhs, rhs)
make_ols_estimates!(b_hat, @view(G_hat_1[:,Not(j)]), @view(B[:,Not(j)]), @view(B[:,j]),K, j, lhs, rhs)
pos = quantile(vec(b_hat), p)
replace!(x -> x>pos ? zero(Float64) : x, b_hat)
replace!(x -> xā¤pos ? 1.0 : x, b_hat)
G_hat_1[Not(j),j] .= b_hat
return nothing
end
sweep_matrix! = function(G_hat_1,B,K, p, n, lhs, rhs)
for j in 1:n
b_hat = Vector{Float64}(undef, n-1) # allocate vector for results
sweep_column_two!(G_hat_1, B, @view(K[Not(j),Not(j)]), j, p, b_hat, lhs, rhs)# update column
K[:,j] .= rand(n) # update K
end
return nothing
end
dummy_one_run = function()
n = 100
p = 0.3
gamma = 0.9
maxit = 1
B = make_B(n).mat
G_hat_1 = G_hat(n).mat
K = rand(n,n)
lhs = init_lhs(n)
rhs = init_rhs(n)
for it in 1:maxit
println(it)
sweep_matrix!(G_hat_1, B,K, p, n, lhs, rhs)
end
return nothing
end
@btime dummy_one_run()
For some reason, I end up creating a huge number of allocations (392480) somewhere in the inner loops when trying to mutate. The allocations make the code run very slowly. Yet when @btime the individual lines within the update_lhs!, update_rhs! functions, they seem to perform very few allocations. I would really appreciate some help, as I cannot work out where the allocations are happeningā¦