gctl/lib/math/fft.cpp
2025-04-23 12:39:44 +08:00

334 lines
8.3 KiB
C++

/********************************************************
* ██████╗ ██████╗████████╗██╗
* ██╔════╝ ██╔════╝╚══██╔══╝██║
* ██║ ███╗██║ ██║ ██║
* ██║ ██║██║ ██║ ██║
* ╚██████╔╝╚██████╗ ██║ ███████╗
* ╚═════╝ ╚═════╝ ╚═╝ ╚══════╝
* Geophysical Computational Tools & Library (GCTL)
*
* Copyright (c) 2023 Yi Zhang (yizhang-geo@zju.edu.cn)
*
* GCTL is distributed under a dual licensing scheme. You can redistribute
* it and/or modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation, either version 2
* of the License, or (at your option) any later version. You should have
* received a copy of the GNU Lesser General Public License along with this
* program. If not, see <http://www.gnu.org/licenses/>.
*
* If the terms and conditions of the LGPL v.2. would prevent you from using
* the GCTL, please consider the option to obtain a commercial license for a
* fee. These licenses are offered by the GCTL's original author. As a rule,
* licenses are provided "as-is", unlimited in time for a one time fee. Please
* send corresponding requests to: yizhang-geo@zju.edu.cn. Please do not forget
* to include some description of your company and the realm of its activities.
* Also add information on how to contact you by electronic and paper mail.
******************************************************/
#include "fft.h"
#ifdef GCTL_FFTW3
void gctl::dft_r2c_1d(const _1d_array &in_real, _1cd_array &out_spectrum, double sampling, _1d_array *freq_ptr)
{
if (in_real.empty())
{
throw std::invalid_argument("[GCTL] Invalid input array size for gctl::dft_r2c_1d(...)");
}
int m = in_real.size();
int n = m/2 + 1; // out size
double *in = fftw_alloc_real(m);
fftw_complex *out= fftw_alloc_complex(n);
fftw_plan p = fftw_plan_dft_r2c_1d(m, in, out, FFTW_ESTIMATE);
for (int i = 0; i < m; i++)
{
in[i] = in_real[i];
}
fftw_execute(p);
fftw_destroy_plan(p);
out_spectrum.resize(n);
for (int i = 0; i < n; i++)
{
// 这里我们把2.0改成了1.0
out_spectrum[i].real(out[i][0]*2.0/m);
out_spectrum[i].imag(out[i][1]*2.0/m);
}
if (freq_ptr != nullptr)
{
freq_ptr->resize(n);
for (size_t i = 0; i < n; i++)
{
freq_ptr->at(i) = 0.5*sampling*i/n;
}
}
fftw_free(in);
fftw_free(out);
return;
}
void gctl::dft_c2r_1d(const _1cd_array &in_spectrum, _1d_array &out_real)
{
if (in_spectrum.empty())
{
throw std::invalid_argument("[GCTL] Invalid input array size for gctl::dft_c2r_1d(...)");
}
int m = in_spectrum.size();
int n = 2*(m-1);
fftw_complex *in = fftw_alloc_complex(m);
double *out= fftw_alloc_real(n);
fftw_plan p = fftw_plan_dft_c2r_1d(n, in, out, FFTW_ESTIMATE);
for (int i = 0; i < m; i++)
{
// 这里我们删除了 *0.5
in[i][0] = in_spectrum[i].real()*0.5*n;
in[i][1] = in_spectrum[i].imag()*0.5*n;
}
fftw_execute(p);
fftw_destroy_plan(p);
out_real.resize(n);
for (int i = 0; i < n; i++)
{
out_real[i] = out[i]/n;
}
fftw_free(in);
fftw_free(out);
return;
}
void gctl::dft2d(const _2d_matrix &in_arr, _1cd_array &out_arr)
{
if (in_arr.empty())
{
throw length_error("The input array is empty. From gctl::dft2d(...)");
}
int m = in_arr.row_size();
int n = in_arr.col_size();
fftw_complex *in = (fftw_complex*) fftw_malloc(sizeof(fftw_complex)*m*n);
fftw_complex *out= (fftw_complex*) fftw_malloc(sizeof(fftw_complex)*m*n);
for (int i = 0; i < m; i++)
{
for (int j = 0; j < n; j++)
{
in[i*n+j][0] = in_arr[i][j];
in[i*n+j][1] = 0.0;
}
}
fftw_plan p = fftw_plan_dft_2d(m, n, in, out, FFTW_FORWARD, FFTW_ESTIMATE);
fftw_execute(p);
out_arr.resize(m*n);
for (int i = 0; i < m*n; i++)
{
out_arr[i].real(out[i][0]);
out_arr[i].imag(out[i][1]);
}
fftw_destroy_plan(p);
fftw_free(in);
fftw_free(out);
return;
}
void gctl::idft2d(const _1cd_array &in_arr, int m, int n, _2d_matrix &out_arr)
{
if (in_arr.empty())
{
throw length_error("The input array is empty. From gctl::idft2d(...)");
}
if (m <= 0 || n <= 0 || m*n != in_arr.size())
{
throw invalid_argument("Invalid parameter. From gctl::idft2d(...)");
}
fftw_complex *in = (fftw_complex*) fftw_malloc(sizeof(fftw_complex)*m*n);
fftw_complex *out= (fftw_complex*) fftw_malloc(sizeof(fftw_complex)*m*n);
for (int i = 0; i < m*n; i++)
{
in[i][0] = in_arr[i].real();
in[i][1] = in_arr[i].imag();
}
fftw_plan p = fftw_plan_dft_2d(m, n, in, out, FFTW_BACKWARD, FFTW_ESTIMATE);
fftw_execute(p);
out_arr.resize(m, n);
for (int i = 0; i < m; i++)
{
for (int j = 0; j < n; j++)
{
out_arr[i][j] = out[i*n+j][0]/(m*n);
}
}
fftw_destroy_plan(p);
fftw_free(in);
fftw_free(out);
return;
}
void gctl::dct1d(const _1d_array &in_arr, _1d_array &out_arr)
{
if (in_arr.empty())
{
throw length_error("The input array is empty. From gctl::dct1d(...)");
}
int m = in_arr.size();
out_arr.resize(m);
fftw_plan p = fftw_plan_r2r_1d(m, in_arr.get(), out_arr.get(),
FFTW_REDFT10, FFTW_ESTIMATE);
fftw_execute(p);
fftw_destroy_plan(p);
return;
}
void gctl::idct1d(const _1d_array &in_arr, _1d_array &out_arr)
{
if (in_arr.empty())
{
throw length_error("The input array is empty. From gctl::idct1d(...)");
}
int m = in_arr.size();
out_arr.resize(m);
fftw_plan p = fftw_plan_r2r_1d(m, in_arr.get(), out_arr.get(),
FFTW_REDFT01, FFTW_ESTIMATE);
fftw_execute(p);
for (int i = 0; i < m; i++)
{
out_arr[i] = out_arr[i]*0.5/m;
}
fftw_destroy_plan(p);
return;
}
#else
void gctl::dft_r2c_1d(const _1d_array &in_real, _1cd_array &out_spectrum, double sampling, _1d_array *freq_ptr)
{
if (in_real.empty())
{
throw std::invalid_argument("[GCTL] Invalid input array size for gctl::dft_r2c_1d(...)");
}
int n = in_real.size();
int s = n/2 + 1; // out size
int m = log2(n);
double *ar, *ai;
ar = new double [n];
ai = new double [n];
for (size_t i = 0; i < n; i++)
{
ar[i] = in_real[i];
ai[i] = 0.0;
}
int i,j,k,k1,l1,l2,ip;
double *ir,*ii,*i1,*i2;
double f1,t1,t2,v,u,w1,w2,d,vu,pi1;
f1=(double)1;
j=0;
for(i=0,ir=ar,ii=ai;i<n;i++,ir++,ii++)
{ if(i<j)
{ i1=ar+j; i2=ai+j;
t1=*i1; t2=*i2; *i1=*ir; *i2=*ii; *ir=t1; *ii=t2;
}
k1=n/2;
while((j>=k1)&&(k1!=0))
{ j=j-k1; k1=k1/2;
}
j=j+k1;
}
for(j=1;j<m+1;j++)
{ l1=2;
for (i=1;i<j;i++) l1=2*l1;
l2=l1/2; pi1=(double)GCTL_Pi/(double)l2; v=(double)1; u=(double)0;
w1=(double)cos(pi1); w2=(double)sin(pi1);
for(k=1;k<l2+1;k++)
{ d=f1*u;
for(i=k-1;i<n-l1+k;i+=l1)
{ ip=i+l2; i1=ar+ip; i2=ai+ip; ir=ar+i; ii=ai+i;
t1=*i1*v+*i2*d; t2=*i2*v-*i1*d;
*i1=*ir-t1; *i2=*ii-t2; *ir=*ir+t1; *ii=*ii+t2;
}
vu=v*w1-u*w2; u=v*w2+u*w1; v=vu;
} }
for(i=0,ir=ar,ii=ai;i<n;i++,ir++,ii++)
{
*ir/=n; *ii/=n;
}
out_spectrum.resize(s);
for (int i = 0; i < s; i++)
{
out_spectrum[i].real(ar[i]);
out_spectrum[i].imag(ai[i]);
}
if (freq_ptr != nullptr)
{
freq_ptr->resize(s);
for (size_t i = 0; i < s; i++)
{
freq_ptr->at(i) = 0.5*sampling*i/s;
}
}
delete[] ar;
delete[] ai;
return;
}
void gctl::dft_c2r_1d(const _1cd_array &in_spectrum, _1d_array &out_real)
{
return;
}
void gctl::dft2d(const _2d_matrix &in_arr, _1cd_array &out_arr)
{
return;
}
void gctl::idft2d(const _1cd_array &in_arr, int m, int n, _2d_matrix &out_arr)
{
return;
}
void gctl::dct1d(const _1d_array &in_arr, _1d_array &out_arr)
{
return;
}
void gctl::idct1d(const _1d_array &in_arr, _1d_array &out_arr)
{
return;
}
#endif // GCTL_FFTW3