Fast Fourier Transform for tensor product spaces in Julia

user_2820579 0
Viewed 128

Right now I use the package FFTW in order to get some Fourier Transforms I am interested in. However, I'm wondering if there is already a package of FFT that can do the transformation in a vector space which is of the form kron(C2, Rn), where C2 means a 2x2 system and Rn represents the "spatial" subspace in which one is interested in getting the Fourier Transform. In other words, does it exist a routine that implements:

kron(Id2x2, FFT)[kron(C2, Rn)] = kron(C2, FFT(Rn))

Of course the real problem I am interested is in the "two particle case" where the vector space (Hilbert space) is kron(kron(C2, Rn),kron(C2, Rn)), so in this case the routine would need an operator like kron(kron(Id2x2, FFT), kron(Id2x2, FFT)).

Note 1: I haven't tried to do the problem taking partial traces, but in my case this option simply may not work because the states are sparse, i.e. it might be ineficient.

Note 2: Note that (unless I'm mistaken) for kron(C2, Rn) one could do "twice" the fft (one in each sector of C2). However this might also be ineficient for large vector spaces.

julia physics fft
1 Answers

Here's an example of what I think you are asking. res is computed by FFT from mat = kron(C2, Rn), and this is (as you say) a wasteful way of doing kron(C2, fft(Rn)) since it the FFT along the k dimension is re-done for each of the 2×2 other dimensions. But the point, presumably, is to do this for "entangled" states in the product space -- a generic likemat = rand(8,2) cannot be decomposed into factors kron(likeC2, likeRn).

(If instead you are really only interested in "un-entangled" product states, then you should probably just work with their components. Combining with kron will then always be wasteful. The package Kronecker.jl may help for some things, but I don't think it knows about fft.)

This uses my package to handle kron-like operations; you could just write out the necessary reshapes yourself, too.

julia> C2 = [1 2; 3 4]; Rn = [1,10,0,0];

julia> mat = kron(C2,Rn)
8×2 Matrix{Int64}:
  1   2
 10  20
  0   0
  0   0
  3   4
 30  40
  0   0
  0   0

julia> using TensorCast, FFTW

# notation: kron is a reshape of a tensor product, to combine i & k

julia> kron(C2,Rn) == @cast out[(k,i),j] := C2[i,j] * Rn[k]
true

# reshape mat to put the index from Rn in its own dimension:

julia> @cast tri[k,i,j] := mat[(k,i),j]  (i in 1:2);

julia> summary(tri)
"4×2×2 Array{Int64, 3}"

# then fft(tri, 1) is the FFT along only that, reshape back:

julia> @cast res[(ktil,i),j] := fft(tri, 1)[ktil,i,j]
8×2 Matrix{ComplexF64}:
  11.0+0.0im    22.0+0.0im
   1.0-10.0im    2.0-20.0im
  -9.0+0.0im   -18.0+0.0im
   1.0+10.0im    2.0+20.0im
  33.0+0.0im    44.0+0.0im
   3.0-30.0im    4.0-40.0im
 -27.0+0.0im   -36.0+0.0im
   3.0+30.0im    4.0+40.0im

julia> res ≈ kron(C2, fft(Rn))
true

julia> res ≈ fft(mat, 1)
false

julia> fft(Rn)
4-element Vector{ComplexF64}:
 11.0 + 0.0im
  1.0 - 10.0im
 -9.0 + 0.0im
  1.0 + 10.0im

# if fft() understood the dims keyword, it could be tidier:

julia> _fft(x; dims) = fft(x, dims);

julia> @cast _res[(k,i),j] := _fft(k) mat[(k,i),j]  (i in 1:2);

julia> _res ≈ res
true
user_137067350
Related