Generating runnable code from symbolic expressions

Here we show a small example of how to translate a derived expression into runnable code. For this we have the print_julia_function function which translates an expression into a julia function that can be run using the TensorOperations.jl package.

As an example we load the omega_aibj_r expression from the CCSD example.

julia> using SpinAdaptedSecondQuantization
julia> using Serialization
julia> omega_aibj_r = deserialize("omega_aibj_r")∑_c(F_ac t_bjci)
- ∑_k(F_ki t_akbj)
- ∑_ck(g_acki t_bjck)
- ∑_ck(g_ackj t_bkci)
+ ∑_kc(g_aikc u_bjck)
- ∑_kcld(g_kcld t_aibk u_cjdl)
- ∑_kcld(g_kcld t_aicj u_bkdl)

Then we can turn this into a function:

julia> print_julia_function("omega_doubles_r", omega_aibj_r);
julia> println(ans)function omega_doubles_r(no, nv, F, t, g, u)
    nao = no + nv
    o = 1:no
    v = no+1:nao
    
    X = zeros(nv, no, nv, no)
    g_vvoo = @view g[v,v,o,o]
    g_voov = @view g[v,o,o,v]
    g_ovov = @view g[o,v,o,v]
    t_vovo = @view t[v,o,v,o]
    u_vovo = @view u[v,o,v,o]
    F_vv = @view F[v,v]
    F_oo = @view F[o,o]
    @tensoropt (a=>10χ,b=>10χ,c=>10χ,d=>10χ,i=>χ,j=>χ,k=>χ,l=>χ) begin
        X[a,i,b,j] += F_vv[a,c]*t_vovo[b,j,c,i]
        X[a,i,b,j] += -F_oo[k,i]*t_vovo[a,k,b,j]
        X[a,i,b,j] += -g_vvoo[a,c,k,i]*t_vovo[b,j,c,k]
        X[a,i,b,j] += -g_vvoo[a,c,k,j]*t_vovo[b,k,c,i]
        X[a,i,b,j] += g_voov[a,i,k,c]*u_vovo[b,j,c,k]
        X[a,i,b,j] += -g_ovov[k,c,l,d]*t_vovo[a,i,b,k]*u_vovo[c,j,d,l]
        X[a,i,b,j] += -g_ovov[k,c,l,d]*t_vovo[a,i,c,j]*u_vovo[b,k,d,l]
    end
    X
end

By default the generated function will take in tensors with all indices assumed to run over all molecular orbitals and make views into these to get sub-blocks where some indices run over only occupied or virtual indices. For integral tensors this can be nice if many different blocks are used, then the caller of the function does not have to think about which blocks to give to the function. However, for tensors such as the T amplitudes where only one specific block exists (the ovov...) block, it is better to limit the tensor to only these indices, rather than embedding it in a much larger tensor. To achieve this one can specify the optional keyword argument explicit_tensor_blocks which makes all sub-blocks used of these tensors be specified as individual inputs to the generated function.

julia> print_julia_function("omega_doubles_r", omega_aibj_r;
           explicit_tensor_blocks = ["t", "u"]);
julia> println(ans)function omega_doubles_r(no, nv, F, t_vovo, g, u_vovo)
    nao = no + nv
    o = 1:no
    v = no+1:nao
    
    X = zeros(nv, no, nv, no)
    g_vvoo = @view g[v,v,o,o]
    g_voov = @view g[v,o,o,v]
    g_ovov = @view g[o,v,o,v]
    F_vv = @view F[v,v]
    F_oo = @view F[o,o]
    @tensoropt (a=>10χ,b=>10χ,c=>10χ,d=>10χ,i=>χ,j=>χ,k=>χ,l=>χ) begin
        X[a,i,b,j] += F_vv[a,c]*t_vovo[b,j,c,i]
        X[a,i,b,j] += -F_oo[k,i]*t_vovo[a,k,b,j]
        X[a,i,b,j] += -g_vvoo[a,c,k,i]*t_vovo[b,j,c,k]
        X[a,i,b,j] += -g_vvoo[a,c,k,j]*t_vovo[b,k,c,i]
        X[a,i,b,j] += g_voov[a,i,k,c]*u_vovo[b,j,c,k]
        X[a,i,b,j] += -g_ovov[k,c,l,d]*t_vovo[a,i,b,k]*u_vovo[c,j,d,l]
        X[a,i,b,j] += -g_ovov[k,c,l,d]*t_vovo[a,i,c,j]*u_vovo[b,k,d,l]
    end
    X
end

After generating the function, it can be copied into a new julia session

julia> using TensorOperations
julia> function omega_doubles_r(no, nv, F, t_vovo, g, u_vovo)
           nao = no + nv
           o = 1:no
           v = no+1:nao
       
           X = zeros(nv, no, nv, no)
           g_vvoo = @view g[v,v,o,o]
           g_voov = @view g[v,o,o,v]
           g_ovov = @view g[o,v,o,v]
           F_vv = @view F[v,v]
           F_oo = @view F[o,o]
           @tensoropt (a=>10χ,b=>10χ,c=>10χ,d=>10χ,i=>χ,j=>χ,k=>χ,l=>χ) begin
               X[a,i,b,j] += F_vv[a,c]*t_vovo[b,j,c,i]
               X[a,i,b,j] += -F_oo[k,i]*t_vovo[a,k,b,j]
               X[a,i,b,j] += -g_vvoo[a,c,k,i]*t_vovo[b,j,c,k]
               X[a,i,b,j] += -g_vvoo[a,c,k,j]*t_vovo[b,k,c,i]
               X[a,i,b,j] += g_voov[a,i,k,c]*u_vovo[b,j,c,k]
               X[a,i,b,j] += -g_ovov[k,c,l,d]*t_vovo[a,i,b,k]*u_vovo[c,j,d,l]
               X[a,i,b,j] += -g_ovov[k,c,l,d]*t_vovo[a,i,c,j]*u_vovo[b,k,d,l]
           end
           X
       endomega_doubles_r (generic function with 1 method)

Then to test it we can setup some test tensors.

julia> no = 4 # Number of occupied orbitals4
julia> nv = 10 # Number of virtual orbitals10
julia> nmo = no + nv # Number of molecular orbitals14
julia> F = randn(nmo, nmo) # Fock matrix populated by random numbers14×14 Matrix{Float64}:
  0.188046   -0.21162    -0.0386384  …  -0.45207     0.367636   0.434867
 -0.253713    2.8659     -1.24068       -0.984165    0.8137     0.0649491
  0.421622   -1.06356     0.979118      -1.44017    -0.027073  -2.46412
 -0.629281    1.97271    -0.0338134     -1.40376     0.414923  -1.7002
 -1.61752    -0.742114   -0.114975       1.13946     1.04346    0.658384
 -0.136062    0.785876    0.068434   …   0.618974   -2.22138   -0.213194
 -0.221601    0.804374    1.12313        0.296344    0.703773  -2.35013
 -0.0378854  -0.714178    0.380671      -0.679157    0.768667  -0.0700148
 -1.669       0.307124   -0.0762624      1.36232    -0.759661   0.36041
  0.414556   -2.16142    -0.174832      -0.0187662   3.50552    1.50287
  0.0526091   2.06127    -2.53096    …  -0.694197   -0.767314   1.78049
  0.290051    0.374514   -1.28515        0.907994   -1.25619   -1.50066
  0.10718     0.725693    0.204246      -2.19833    -0.294563   0.494499
  0.220802    0.0469061  -0.548947       0.0846749   0.101961   0.670323
julia> g = randn(nmo, nmo, nmo, nmo);
julia> t = randn(nv, no, nv, no);
julia> u = 2 * t - PermutedDimsArray(t, (1, 4, 3, 2));

Then we can call the function and post-process the output by symmetrizing and scaling the diagonal

julia> Ω_aibj = @time omega_doubles_r(no, nv, F, t, g, u);  6.761660 seconds (18.06 M allocations: 900.395 MiB, 4.91% gc time, 99.99% compilation time)
julia> Ω_aibj += PermutedDimsArray(t, (3, 4, 1, 2)); # Symmetrize
julia> for a in 1:nv, i in 1:no
           Ω_aibj[a, i, a, i] *= 0.5 # Scale diagonal by 0.5
       end