Minimize memory allocations while calculating the matrix representation of a Hamiltonian

Question

I am writing a script using Julia (calling my functions from the REPL) to compute the matrix representation of certain Hamiltonian. After I run my script I see that the memory allocations count is considerably high:

julia> @time hamiltonian(40, 20);
      5.860412 seconds (76.33 M allocations: 2.085 GiB, 23.78% gc time)

The time it takes to execute has not been an issue, but I would like to know how could I write this script as a more idiomatic (but still optimal) Julia code while trying to minimize the memory allocations?

This is my script until now:

const ω = ω0 = 1
const γ = 1

myFactorial(x) = factorial(big(x))
G(j) = (2 * γ) / (ω * sqrt(2*j))
    
function cMinus(j::Int, m::Real)
    sqrt(j * (j + 1) - m * (m - 1))
end
    
function cPlus(j::Int, m::Real)
    sqrt(j * (j + 1) - m * (m + 1))
end


# Function to compute the overlap in the BCE base.
function overlap(
    n_bra::Int,
    m_bra::Real,
    n_ket::Int,
    m_ket::Real,
    j::Int
    )
    newG(m_bra, m_ket) = G(j) * (m_bra - m_ket)
    sum_num(n_bra, n_ket, k) = (-1)^(n_ket - k) * 
        sqrt(bigFactorial(n_bra) * bigFactorial(n_ket)) * 
        newG(m_bra, m_ket)^(n_bra + n_ket -2*k)
    sum_den(n_bra, n_ket, k) = bigFactorial(k) * 
        bigFactorial(n_bra - k) * 
        bigFactorial(n_ket - k)
    if m_bra == m_ket && n_bra != n_ket
        return 0
    else
        term1 = exp(- newG(m_bra, m_ket)^2 / 2)
        # term2 = [sum_num(n_bra, n_ket, k) / sum_den(n_bra, n_ket, k) 
        #   for k=0:min(n_bra, n_ket)]
        term2 = sum(sum_num(n_bra, n_ket, k) / 
            sum_den(n_bra, n_ket, k) 
            for k=0:min(n_bra, n_ket))
        return term1 * sum(term2)
    end
end


# Matrix representation of the Hamiltonian for the specified `n_max` and `j`.
function hamiltonian(n_max::Int, j::Int)
    n_range = 0:n_max
    m_range = -j:j
    dim = (2*j + 1) * (n_max + 1)
    temp = zeros(dim, dim)
    col = 1
    for n_ket = n_range, m_ket = m_range
        row = 1
        for n_bra = n_range, m_bra = m_range
            if row >= col
                # Me encuentro en el estado (n_bra, m_bra, n_ket, m_ket)
                if n_bra == n_ket && m_bra == m_ket
                    temp[row, col] = ω * (n_ket - (G(j) * m_ket)^2)
                elseif m_bra == m_ket + 1
                    temp[row, col] = temp[col, row] = (- ω0 / 2) * 
                        cPlus(j, m_ket) * 
                        overlap(n_bra, m_bra, n_ket, m_ket, j)
                elseif m_bra == m_ket - 1
                    temp[row, col] = temp[col, row] = (- ω0 / 2) * 
                        cMinus(j, m_ket) * 
                        overlap(n_bra, m_bra, n_ket, m_ket, j)
                end
            end
            row += 1
        end
        col += 1
    end
    temp
end

One minimal example of my script working is the following:

julia> @time hamiltonian(1, 1)
  0.000064 seconds (881 allocations: 27.984 KiB)
6×6 Matrix{Float64}:
 -2.0       -0.26013    0.0        0.0       -0.367879   0.0
 -0.26013    0.0       -0.26013    0.367879   0.0       -0.367879
  0.0       -0.26013   -2.0        0.0        0.367879   0.0
  0.0        0.367879   0.0       -1.0        0.26013    0.0
 -0.367879   0.0        0.367879   0.26013    1.0        0.26013
  0.0       -0.367879   0.0        0.0        0.26013   -1.0

I was able to reduce the memory allocations by considering the symmetry of the problem. The approach I took was to compute the values for the main diagonal and those located in the lower-half of the matrix. Then those values are "transposed" to fill the upper-half as M[i,j] = M[j,i]. An example of the time and allocations reduction is shown below:

julia> @time hamiltonian(40, 20)
  2.917963 seconds (39.69 M allocations: 1.129 GiB, 25.57% gc time)

Although, I would like to call hamiltonian(n_max, j) with much higher values of n_max and j, around n_max=300 and j=200.

---

Some mathematical background

For sake of completeness the mathematical expression of the Hamiltonian is:

$$ \begin{align} \langle N', m'_{x} | \hat{H} | N, m_{x} \rangle = & \omega \left[N - (G m_{x})^{2}\right] \delta_{N' N} \delta_{m'_{x} m_{x}} \\ & - \frac{\omega_{0}}{2} \left[C_{+}(j,m_{x}) \delta_{m'_{x} (m_{x}+1)} + C_{-}(j,m_{x}) \delta_{m'_{x} (m_{x}-1)}\right] {}_{m_{x}}\langle N' | N \rangle_{m_{x}} \end{align} $$

The last term in \$\langle\hat{H}\rangle\$ is the overlap between states, it is given by:

$$ \begin{align} {}_{m'_{x}}\langle N' | N \rangle_{m_{x}} = & \exp\left[-\frac{G^{2}}{2} (m'_{x} - m_{x})^2\right] \\ & \sum_{k=0}^{min\{N',N\}} \frac{(-1)^{N-k} \sqrt{N! N'!}}{k! (N-k)! (N'-k)!} \left[G(m'_{x} - m_{x})\right]^{N+N'-2k} \end{align} $$

Additionally, the \$C\$ coefficients are obtained as:

$$ C_{\pm}(j,m) = \sqrt{j(j + 1) - m(m \pm 1)} $$

The following constants are used in this specific case:

$$ \omega = \omega_{0} = \gamma = 1 $$

And the constant \$G\$ is:

$$ G = \frac{2 \gamma}{\omega \sqrt{2 j}}$$

Answer 1

The first problem is you want const ω = const ω0 = 1. The second problem is harder, but basically you are computing way too many factorials rather than computing sum_num(n_bra, n_ket, k) / sum_den(n_bra, n_ket, k), for each k you should look for the formula that computes the k+1th term from the kth term. This will only use multiplication/division, so you won't need to use factorials. If you don't need factorials, you won't get overflow, so you won't need BigFloats and your allocations will vanish, and your code will be about 100x faster.