void fftn_init(int rank, int *n)
/*< initialize >*/
{
int i;
num=1;
for(i=0; i<rank; i++) num*=n[i];
tmp=(fftwf_complex*)fftwf_malloc(sizeof(fftwf_complex)*num);
fftn=fftwf_plan_dft(rank, n, tmp, tmp, FFTW_FORWARD, FFTW_MEASURE);
ifftn=fftwf_plan_dft(rank, n, tmp, tmp, FFTW_BACKWARD, FFTW_MEASURE);
}
static PlanType create(const std::array<std::size_t,NDims>& _shape,
ComplexType* _in,
ComplexType* _out,
fftw_direction _dir = fftw_direction::forward,
unsigned plan_flags = FFTW_MEASURE){
std::array<int,NDims> converted;
for(int i = 0;i < NDims;++i)
converted[i] = _shape[i];
PlanType value = fftwf_plan_dft(NDims,
converted.data(),
_in,
_out,
static_cast<int>(_dir),
plan_flags );
return value;
}
int shrinkWrap
(
float * const & rIntensity,
const std::vector<unsigned> & rSize,
unsigned rnCycles,
float rTargetError,
float rHioBeta,
float rIntensityCutOffAutoCorel,
float rIntensityCutOff,
float rSigma0,
float rSigmaChange,
unsigned rnHioCycles
)
{
if ( rSize.size() != 2 ) return 1;
const unsigned & Ny = rSize[1];
const unsigned & Nx = rSize[0];
/* Evaluate input parameters and fill with default values if necessary */
if ( rIntensity == NULL ) return 1;
if ( rTargetError <= 0 ) rTargetError = 1e-5;
if ( rnHioCycles == 0 ) rnHioCycles = 20;
if ( rHioBeta <= 0 ) rHioBeta = 0.9;
if ( rIntensityCutOffAutoCorel <= 0 ) rIntensityCutOffAutoCorel = 0.04;
if ( rIntensityCutOff <= 0 ) rIntensityCutOff = 0.2;
if ( rSigma0 <= 0 ) rSigma0 = 3.0;
if ( rSigmaChange <= 0 ) rSigmaChange = 0.01;
float sigma = rSigma0;
/* calculate this (length of array) often needed value */
unsigned nElements = 1;
for ( unsigned i = 0; i < rSize.size(); ++i )
{
assert( rSize[i] > 0 );
nElements *= rSize[i];
}
/* allocate needed memory so that HIO doesn't need to allocate and
* deallocate on each call */
fftwf_complex * const curData = fftwf_alloc_complex( nElements );
fftwf_complex * const gPrevious = fftwf_alloc_complex( nElements );
auto const isMasked = new float[nElements];
/* create fft plans G' to g' and g to G */
auto toRealSpace = fftwf_plan_dft( rSize.size(),
(int*) &rSize[0], curData, curData, FFTW_BACKWARD, FFTW_ESTIMATE );
auto toFreqSpace = fftwf_plan_dft( rSize.size(),
(int*) &rSize[0], gPrevious, curData, FFTW_FORWARD, FFTW_ESTIMATE );
/* create first guess for mask from autocorrelation (fourier transform
* of the intensity @see
* https://en.wikipedia.org/wiki/Wiener%E2%80%93Khinchin_theorem */
#pragma omp parallel for
for ( unsigned i = 0; i < nElements; ++i )
{
curData[i][0] = rIntensity[i]; /* Re */
curData[i][1] = 0;
}
fftwf_execute( toRealSpace );
complexNormElementwise( isMasked, curData, nElements );
/* fftShift is not necessary, but I introduced this, because for the
* example it shifted the result to a better looking position ... */
//fftShift( isMasked, Nx,Ny );
libs::gaussianBlur( isMasked, Nx, Ny, sigma );
#if DEBUG_SHRINKWRAPP_CPP == 1
std::ofstream file;
std::string fname = std::string("shrinkWrap-init-mask-blurred");
file.open( ( fname + std::string(".dat") ).c_str() );
for ( unsigned ix = 0; ix < rSize[0]; ++ix )
{
for ( unsigned iy = 0; iy < rSize[1]; ++iy )
file << std::setw(10) << isMasked[ iy*rSize[0] + ix ] << " ";
file << "\n";
}
file.close();
std::cout << "Written out " << fname << ".png\n";
#endif
/* apply threshold to make binary mask */
{
const auto absMax = vectorMax( isMasked, nElements );
const float threshold = rIntensityCutOffAutoCorel * absMax;
#pragma omp parallel for
for ( unsigned i = 0; i < nElements; ++i )
isMasked[i] = isMasked[i] < threshold ? 1 : 0;
}
#if DEBUG_SHRINKWRAPP_CPP == 1
fname = std::string("shrinkWrap-init-mask");
file.open( ( fname + std::string(".dat") ).c_str() );
for ( unsigned ix = 0; ix < rSize[0]; ++ix )
{
for ( unsigned iy = 0; iy < rSize[1]; ++iy )
file << std::setw(10) << isMasked[ iy*rSize[0] + ix ] << " ";
file << "\n";
}
file.close();
std::cout << "Written out " << fname << ".png\n";
#endif
/* copy original image into fftw_complex array and add random phase */
#pragma omp parallel for
for ( unsigned i = 0; i < nElements; ++i )
{
curData[i][0] = rIntensity[i]; /* Re */
curData[i][1] = 0;
}
/* in the first step the last value for g is to be approximated
* by g'. The last value for g, called g_k is needed, because
* g_{k+1} = g_k - hioBeta * g' ! This is inside the loop
* because the fft is needed */
#pragma omp parallel for
for ( unsigned i = 0; i < nElements; ++i )
{
gPrevious[i][0] = curData[i][0];
gPrevious[i][1] = curData[i][1];
}
/* repeatedly call HIO algorithm and change mask */
for ( unsigned iCycleShrinkWrap = 0; iCycleShrinkWrap < rnCycles; ++iCycleShrinkWrap )
{
/************************** Update Mask ***************************/
std::cout << "Update Mask with sigma=" << sigma << "\n";
/* blur |g'| (normally g' should be real!, so |.| not necessary) */
complexNormElementwise( isMasked, curData, nElements );
libs::gaussianBlur( isMasked, Nx, Ny, sigma );
const auto absMax = vectorMax( isMasked, nElements );
/* apply threshold to make binary mask */
const float threshold = rIntensityCutOff * absMax;
#pragma omp parallel for
for ( unsigned i = 0; i < nElements; ++i )
isMasked[i] = isMasked[i] < threshold ? 1 : 0;
/* update the blurring sigma */
sigma = fmax( 1.5, ( 1 - rSigmaChange ) * sigma );
for ( unsigned iHioCycle = 0; iHioCycle < rnHioCycles; ++iHioCycle )
{
/* apply domain constraints to g' to get g */
#pragma omp parallel for
for ( unsigned i = 0; i < nElements; ++i )
{
if ( isMasked[i] == 1 or /* g' */ curData[i][0] < 0 )
{
gPrevious[i][0] -= rHioBeta * curData[i][0];
gPrevious[i][1] -= rHioBeta * curData[i][1];
}
else
{
gPrevious[i][0] = curData[i][0];
gPrevious[i][1] = curData[i][1];
}
}
/* Transform new guess g for f back into frequency space G' */
fftwf_execute( toFreqSpace );
/* Replace absolute of G' with measured absolute |F| */
applyComplexModulus( curData, curData, rIntensity, nElements );
fftwf_execute( toRealSpace );
} // HIO loop
/* check if we are done */
const float currentError = imresh::libs::calculateHioError( curData /*g'*/, isMasked, nElements );
std::cout << "[Error " << currentError << "/" << rTargetError << "] "
<< "[Cycle " << iCycleShrinkWrap << "/" << rnCycles-1 << "]"
<< "\n";
if ( rTargetError > 0 && currentError < rTargetError )
break;
if ( iCycleShrinkWrap >= rnCycles )
break;
} // shrink wrap loop
for ( unsigned i = 0; i < nElements; ++i )
rIntensity[i] = curData[i][0];
/* free buffers and plans */
fftwf_destroy_plan( toFreqSpace );
fftwf_destroy_plan( toRealSpace );
fftwf_free( curData );
fftwf_free( gPrevious);
delete[] isMasked;
return 0;
}
int main(int argc, char* argv[])
{
/*define variables*/
int nx,nx1,nt;
int n1,n2;
float d1,o1,d2,o2;
int padt,padx;
int ntfft,*n,nw,nk;
float **d,*wavelet,**shot,**ds,**vel,**vmig,**M,v_ave;
float *kx,*omega,dkx,dw;
sf_complex **m,**ms,**mr,*in2a,*in2b,*cs,*cr,*c,czero;
sf_complex Ls;
float fmin,fmax,f_low,f_high;
int if_low,if_high;
int ix,iw,ik;
float dt,dx,ox,dz,zmax;
fftwf_plan p2a,p2b;
sf_file in,out,velfile,source_wavelet;
int iz,nz;
int ishot,max_num_shot,ig,ng,it,index;
int iswavelet;
/*define sf input output*/
sf_init (argc,argv);
in = sf_input("in");
out = sf_output("out");
velfile = sf_input("velfile");
if (!sf_histint(in,"n1",&n1)) sf_error("No n1= in input");
if (!sf_histfloat(in,"d1",&d1)) sf_error("No d1= in input");
if (!sf_histfloat(in,"o1",&o1)) o1=0.;
if (!sf_histint(in,"n2",&n2)) sf_error("No n2= in vel");
if (!sf_histfloat(in,"d2",&d2)) sf_error("No d2= in input");
if (!sf_histfloat(in,"o2",&o2)) o2=0.;
dt = d1;
dx = d2;
ox = o2;
nx1 = n2;
nt = n1;
if (!sf_histint(velfile,"n1",&nz)) sf_error("No n1= in vel");
if (!sf_histfloat(velfile,"d1",&dz)) sf_error("No n1= in vel");
if (!sf_histint(velfile,"n2",&n2)) sf_error("No n2= in vel");
if (!sf_getint("iswavelet",&iswavelet)) iswavelet = 0;
source_wavelet=sf_input("source_wavelet");
max_num_shot=100;
ng=700;
nx=n2;
padt = 2;
padx = 2;
ntfft = padt*nt;
nw=ntfft/2+1;
nk = padx*nx;
dw = 2*PI/ntfft/dt;
dkx = 2*PI/nk/dx;
sf_putint(out,"n1",nz);
sf_putint(out,"n2",nx);
sf_putfloat(out,"d1",dz);
sf_putstring(out,"label1","z");
sf_putstring(out,"unit1","m");
sf_putstring(out,"title","migrated");
if (!sf_getfloat("fmax",&fmax)) fmax = 0.5/d1; /* max frequency to process */
if (fmax > 0.5/d1) fmax = 0.5/d1;
if (!sf_getfloat("fmin",&fmin)) fmin = 0.1; /* min frequency to process */
if (!sf_getfloat("Zmax",&zmax)) zmax = (nz-1)*dz; /* max Depth to migrate */
/*define axis variables*/
dkx=(float) 2*PI/nk/dx;
dw=(float) 2*PI/ntfft/dt;
/*allocate memory to dynamic arrays*/
d = sf_floatalloc2(nt,nx1);
shot=sf_floatalloc2(nt,ng);
ds=sf_floatalloc2(nt,nx);
vel = sf_floatalloc2(nz,nx);
wavelet=sf_floatalloc(nt);
vmig = sf_floatalloc2(nz,nx);
m = sf_complexalloc2(nw,nx);
ms = sf_complexalloc2(nw,nx);
mr = sf_complexalloc2(nw,nx);
kx= sf_floatalloc (nk);
omega= sf_floatalloc (nw);
in2a = sf_complexalloc(nk);
in2b = sf_complexalloc(nk);
n = sf_intalloc(1);
M= sf_floatalloc2(nz,nx);
c = sf_complexalloc(nx);
cs = sf_complexalloc(nx);
cr = sf_complexalloc(nx);
/*read input files*/
sf_floatread(d[0],nx1*nt,in);
sf_floatread(vel[0],nx*nz,velfile);
/* If there is no wavelet use delta as default
If there is a wavelet use it*/
if (iswavelet==0) {
for (it=0; it<nt; it++) wavelet[it] = 0.0;
wavelet[0]=1;
}
if (iswavelet==1) sf_floatread(wavelet,nt,source_wavelet);
/* This part is important: we need to define the horizontal wavenumber and frequency axes right.*/
dw = 2*PI/ntfft/dt;
dkx = 2*PI/nk/dx;
for (iw=0;iw<nw;iw++){
omega[iw] = dw*iw;
}
for (ik=0;ik<nk;ik++){
if (ik<nk/2) kx[ik] = dkx*ik;
else kx[ik] = -(dkx*nk - dkx*ik);
}
/* Define minimum and maximum frequency index to process*/
f_low = fmin; /* min frequency to process */
f_high = fmax; /* max frequency to process */
if(f_low>0){
if_low = trunc(f_low*dt*ntfft);
}
else{
if_low = 0;
}
if(f_high*dt*ntfft+1<nw){
if_high = trunc(f_high*dt*ntfft)+1;
}
else{
if_high = nw;
}
__real__ czero = 0;
__imag__ czero = 0;
n[0] = nk;
p2a = fftwf_plan_dft(1, n, (fftwf_complex*)in2a, (fftwf_complex*)in2a, FFTW_FORWARD, FFTW_ESTIMATE);
p2b = fftwf_plan_dft(1, n, (fftwf_complex*)in2b, (fftwf_complex*)in2b, FFTW_BACKWARD, FFTW_ESTIMATE);
fftwf_execute(p2a); /* FFT x to k */
fftwf_execute(p2b); /* FFT x to k */
/* Define initial migrated model and source field as zeros*/
for (iz=0; iz<nz; iz++) {
for (ix=0; ix<nx; ix++) M[ix][iz] = 0.0;
}
for (it=0; it<nt; it++) {
for (ix=0; ix<nx; ix++) ds[ix][it] = 0.0;
}
for (iz=0; iz<nz;iz++){
for (ix=0;ix<nx;ix++) vmig[ix][iz]=vel[ix][iz];
}
/* loop over shots*/
for (ishot=0;ishot<max_num_shot;ishot++){
for (ig=0;ig<ng;ig++){
for (it=0; it<nt; it++)
shot[ig][it]=d[ishot*ng+ig][it];
}
for (it=0; it<nt; it++) {
for (ix=0; ix<nx; ix++) ds[ix][it] = 0.0;
}
index=ishot*nx/max_num_shot;
for (it=0; it<nt; it++) ds[index][it]=wavelet[it];
/* apply fourier transform in time direction t-x ---> w-x*/
my_forward_fft(ms,mr,shot,ds,nt,dt,nx,padt);
for (iw=if_low;iw<if_high;iw++){
for (iz=0; iz<nz;iz++){
v_ave=vmig[0][iz];
my_v_ave (v_ave,vmig,iz,nx);
/*Apply phase shift to source side*/
my_phase_shift(ms,czero,iw,iz,omega,kx,nk,nx,v_ave,in2a,in2b,p2a,p2b,dz,0);
for (ix=0;ix<nx;ix++) {
cs[ix]= in2b[ix];
}
/*Apply phase shift to receiver side*/
my_phase_shift(mr,czero,iw,iz,omega,kx,nk,nx,v_ave,in2a,in2b,p2a,p2b,dz,1);
for (ix=0;ix<nx;ix++) {
cr[ix]= in2b[ix];
}
/*Apply split step correction to source and receiver side wavefields*/
my_split_step_correction (ms,cs,vmig,v_ave,iz,dz,iw,dw,nx,0);
my_split_step_correction (mr,cr,vmig,v_ave,iz,dz,iw,dw,nx,1);
/* Apply cross corrolation as an imaging condition*/
for (ix=0;ix<nx;ix++){
__real__ Ls=crealf(ms[ix][iw]);
__imag__ Ls=- cimagf(ms[ix][iw]);
m[ix][iw]=mr[ix][iw]*Ls;
}
/* Update migrated model by stacking*/
for (ix=0;ix<nx;ix++) M[ix][iz]=M[ix][iz]+2*crealf(m[ix][iw]);
}
}
fprintf(stderr,"\r progress = %6.2f%%",(float) 100*(ishot)/(max_num_shot));
}
sf_floatwrite(M[0],nz*nx,out);
fftwf_destroy_plan(p2a);
fftwf_free(in2a);
fftwf_destroy_plan(p2b);
fftwf_free(in2b);
exit (0);
}
