Optimise inverse fourier transform

Author: vitartas

Project.toml

@@ -25,6 +25,7 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Spglib = "f761d5c5-86db-4880-b97f-9680a7cccfb5"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
@@ -53,6 +54,7 @@ SafeTestsets = "0.1"
 Spglib = "1.0.1"
 StaticArrays = "1.9"
 StructArrays = "0.7"
+Tullio = "0.3.8"
 Unitful = "1.25"
 UnitfulAtomic = "1"
 julia = "1.12"

src/Quoll.jl

@@ -6,6 +6,7 @@ using ArgCheck
 using LinearAlgebra
 using StaticArrays
 using OffsetArrays
+using Tullio
 using Dictionaries
 using AutoHashEquals
 using AxisKeys

src/conversions/canonical/bsparse.jl

@@ -34,41 +34,48 @@ end
 
 # TODO: write another method but with shifts (or alternatively write the method with shifts
 # but force the compiler to remove shifts for matching SH convention)
+# TODO: could use out_keydata here instead
 function fourier_transform_data!(out_operator::DenseOperator, in_operator::BSparseOperator, phases_k)
     in_keydata = get_keydata(in_operator)
     in_sparsity = get_sparsity(in_operator)
     out_data = get_data(out_operator)
     out_basisset = get_basisset(out_operator)
-    out_float = get_float(out_operator)
 
     ilocal2iglobal = get_ilocal2iglobal(in_sparsity)
-    atom2offset = get_atom2offset(out_basisset)
+    atom2basis = get_atom2basis(out_basisset)
+
+    # atom2offset = get_atom2offset(out_basisset)
+    # out_float = get_float(out_operator)
 
     for ((iat, jat), in_block) in pairs(in_keydata)
 
         phases_kij = phases_k[ilocal2iglobal[(iat, jat)]]
-        atom2offset_i = atom2offset[iat]
-        atom2offset_j = atom2offset[jat]
 
-        for jb in axes(in_block, 2)
-            jb_dense = jb + atom2offset_j
-            for ib in axes(in_block, 1)
-                ib_dense = ib + atom2offset_i
+        in_matrix = reshape(in_block, size(in_block, 1) * size(in_block, 2), size(in_block, 3))
+        out_data[atom2basis[iat], atom2basis[jat]] = in_matrix * phases_kij
 
-                tmp = zero(out_float)
-                @inbounds for iR in axes(in_block, 3)
-                    tmp += in_block[ib, jb, iR] * phases_kij[iR]
-                end
+        # atom2offset_i = atom2offset[iat]
+        # atom2offset_j = atom2offset[jat]
 
-                out_data[ib_dense, jb_dense] = tmp
-            end
-        end
+        # for jb in axes(in_block, 2)
+        #     jb_dense = jb + atom2offset_j
+        #     for ib in axes(in_block, 1)
+        #         ib_dense = ib + atom2offset_i
+
+        #         tmp = zero(out_float)
+        #         @inbounds for iR in eachindex(phases_kij)
+        #             # tmp += in_block[ib, jb, iR] * phases_kij[iR]
+        #             tmp = muladd(in_block[ib, jb, iR], phases_kij[iR], tmp)
+        #         end
+
+        #         out_data[ib_dense, jb_dense] = tmp
+        #     end
+        # end
 
         # @tullio tmp[ib,jb] := in_block[ib,jb,R] * phases_kij[R]
         # out_data[atom2basis[iat], atom2basis[jat]] = tmp
     end
 
-    # TODO: I could use Hermitian array instead, probably directly in build stage
     if in_sparsity.hermitian
         for jb_dense in 1:size(out_data, 1)
             for ib_dense in 1:jb_dense
@@ -80,30 +87,35 @@ end
 
 # Only a contribution to out_operator because in_operator contains only a single k-point
 function inv_fourier_transform_data!(out_operator::BSparseOperator, in_operator::DenseOperator, phases_k, weight)
-    out_basisset = get_basisset(out_operator)
     out_keydata = get_keydata(out_operator)
     out_sparsity = get_sparsity(out_operator)
+    in_basisset = get_basisset(out_operator)
     in_data = get_data(in_operator)
 
     ilocal2iglobal = get_ilocal2iglobal(out_sparsity)
-    atom2offset = get_atom2offset(out_basisset)
+    atom2basis = get_atom2basis(in_basisset)
+
+    # atom2offset = get_atom2offset(in_basisset)
 
     for ((iat, jat), out_block) in pairs(out_keydata)
 
-        # Assuming phases is a vector here (i.e. for a single k-point)
-        inv_phases_kij = conj(phases_k[ilocal2iglobal[(iat, jat)]])
-        atom2offset_i = atom2offset[iat]
-        atom2offset_j = atom2offset[jat]
+        inv_phases_kij = conj(@view(phases_k[ilocal2iglobal[(iat, jat)]]))
+        in_data_ij = @view(in_data[atom2basis[iat], atom2basis[jat]])
+
+        @tullio out_block[ib,jb,iR] += weight * real(in_data_ij[ib,jb] * inv_phases_kij[iR])
+
+        # atom2offset_i = atom2offset[iat]
+        # atom2offset_j = atom2offset[jat]
 
-        for iR in axes(out_block, 3)
-            for jb in axes(out_block, 2)
-                jb_dense = jb + atom2offset_j
-                for ib in axes(out_block, 1)
-                    ib_dense = ib + atom2offset_i
-                    out_block[ib, jb, iR] += weight * real(in_data[ib_dense, jb_dense] * inv_phases_kij[iR])
-                end
-            end
-        end
+        # for iR in axes(out_block, 3)
+        #     for jb in axes(out_block, 2)
+        #         jb_dense = jb + atom2offset_j
+        #         for ib in axes(out_block, 1)
+        #             ib_dense = ib + atom2offset_i
+        #             out_block[ib, jb, iR] += weight * real(in_data[ib_dense, jb_dense] * inv_phases_kij[iR])
+        #         end
+        #     end
+        # end
     end
 
 end

src/projections/common.jl

@@ -93,7 +93,6 @@ function perform_core_projection(operators, projected_basis, kgrid::KGrid, my_ik
     return v_operators
 end
 
-# TODO: if not views, then sometimes can result in redundant allocations
 function core_valence_partition(operator::DenseOperator, core_mask::BitVector, valence_mask::BitVector)
     O = get_data(operator)
     O₁₁ = view(O, core_mask, core_mask)