73 lines
1.5 KiB
Lua
73 lines
1.5 KiB
Lua
|
-- operations on complex number
|
|||
|
|
complex = {__mt={} }
|
|||
|
|
|
|||
|
|
function complex.new (r, i)
|
|||
|
|
local new={r=r, i=i or 0}
|
|||
|
|
setmetatable(new,complex.__mt)
|
|||
|
|
return new
|
|||
|
|
end
|
|||
|
|
|
|||
|
|
function complex.__mt.__add (c1, c2)
|
|||
|
|
return complex.new(c1.r + c2.r, c1.i + c2.i)
|
|||
|
|
end
|
|||
|
|
|
|||
|
|
function complex.__mt.__sub (c1, c2)
|
|||
|
|
return complex.new(c1.r - c2.r, c1.i - c2.i)
|
|||
|
|
end
|
|||
|
|
|
|||
|
|
function complex.__mt.__mul (c1, c2)
|
|||
|
|
return complex.new(c1.r*c2.r - c1.i*c2.i,
|
|||
|
|
c1.r*c2.i + c1.i*c2.r)
|
|||
|
|
end
|
|||
|
|
|
|||
|
|
function complex.expi (i)
|
|||
|
|
return complex.new(math.cos(i),math.sin(i))
|
|||
|
|
end
|
|||
|
|
|
|||
|
|
function complex.__mt.__tostring(c)
|
|||
|
|
return "("..c.r..","..c.i..")"
|
|||
|
|
end
|
|||
|
|
|
|||
|
|
|
|||
|
|
-- Cooley–Tukey FFT (in-place, divide-and-conquer)
|
|||
|
|
-- Higher memory requirements and redundancy although more intuitive
|
|||
|
|
function fft(vect)
|
|||
|
|
local n=#vect
|
|||
|
|
if n<=1 then return vect end
|
|||
|
|
-- divide
|
|||
|
|
local odd,even={},{}
|
|||
|
|
for i=1,n,2 do
|
|||
|
|
odd[#odd+1]=vect[i]
|
|||
|
|
even[#even+1]=vect[i+1]
|
|||
|
|
end
|
|||
|
|
-- conquer
|
|||
|
|
fft(even);
|
|||
|
|
fft(odd);
|
|||
|
|
-- combine
|
|||
|
|
for k=1,n/2 do
|
|||
|
|
local t=even[k] * complex.expi(-2*math.pi*(k-1)/n)
|
|||
|
|
vect[k] = odd[k] + t;
|
|||
|
|
vect[k+n/2] = odd[k] - t;
|
|||
|
|
end
|
|||
|
|
return vect
|
|||
|
|
end
|
|||
|
|
|
|||
|
|
function toComplex(vectr)
|
|||
|
|
vect={}
|
|||
|
|
for i,r in ipairs(vectr) do
|
|||
|
|
vect[i]=complex.new(r)
|
|||
|
|
end
|
|||
|
|
return vect
|
|||
|
|
end
|
|||
|
|
|
|||
|
|
-- test
|
|||
|
|
data = toComplex{1, 1, 1, 1, 0, 0, 0, 0};
|
|||
|
|
|
|||
|
|
-- this works for old lua versions & luaJIT (depends on version!)
|
|||
|
|
-- print("orig:", unpack(data))
|
|||
|
|
-- print("fft:", unpack(fft(data)))
|
|||
|
|
|
|||
|
|
-- Beginning with Lua 5.2 you have to write
|
|||
|
|
print("orig:", table.unpack(data))
|
|||
|
|
print("fft:", table.unpack(fft(data)))
|