Commit c1a39967 authored by Martin Reinecke's avatar Martin Reinecke
Browse files

test1

parent e3c40bc8
......@@ -659,41 +659,6 @@ template<typename T> class GridderConfig
return res;
}
pyarr_c<T> grid2dirty_wstacking(const pyarr_c<complex<T>> &grid, double w) const
{
checkArray(grid, "grid", {nu, nv});
auto tmp = makeArray<complex<T>>({nu, nv});
auto ptmp = tmp.mutable_data();
pocketfft::c2c({nu,nv},{grid.strides(0),grid.strides(1)},
{tmp.strides(0), tmp.strides(1)}, {0,1}, pocketfft::BACKWARD,
grid.data(), tmp.mutable_data(), T(1), nthreads);
auto res = makeArray<T>({nx_dirty, ny_dirty});
auto pout = res.mutable_data();
{
py::gil_scoped_release release;
#pragma omp parallel num_threads(nthreads)
{
#pragma omp for schedule(static)
for (size_t i=0; i<nx_dirty; ++i)
{
double fx = (i-nx_dirty/2.)*Pixsize_x();
fx *= fx;
for (size_t j=0; j<ny_dirty; ++j)
{
size_t i2 = nu-nx_dirty/2+i;
if (i2>=nu) i2-=nu;
size_t j2 = nv-ny_dirty/2+j;
if (j2>=nv) j2-=nv;
double fy = (j-ny_dirty/2.)*Pixsize_y();
auto ws = wscreen(fx, fy*fy, w, true);
pout[ny_dirty*i + j] = std::real(ptmp[nv*i2+j2]*cfu[i]*cfv[j]*ws);
}
}
}
}
return res;
}
pyarr_c<T> dirty2grid(const pyarr_c<T> &dirty) const
{
checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
......@@ -742,50 +707,14 @@ template<typename T> class GridderConfig
}
return tmp;
}
inline complex<T> wscreen(double x, double y, double w, bool adjoint) const
{
complex<T> wscreen(double x, double y, double w, bool adjoint) const
{
constexpr double pi = 3.141592653589793238462643383279502884197;
constexpr complex<double> imagunit {0,1};
double n = cos(sqrt(x+y));
complex<T> phase = 2*pi*w*imagunit*(n-1);
double n = cos(sqrt(x+y)), xn = 1./n;
double phase = 2*pi*w*(n-1);
if (adjoint) phase *= -1;
return exp(phase)/n;
}
pyarr_c<complex<T>> dirty2grid_wstacking(const pyarr_c<T> &dirty, double w) const
{
checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
auto pdirty = dirty.data();
auto tmp = makeArray<complex<T>>({nu, nv});
auto ptmp = tmp.mutable_data();
pocketfft::stride_t strides{tmp.strides(0),tmp.strides(1)};
{
py::gil_scoped_release release;
for (size_t i=0; i<nu*nv; ++i)
ptmp[i] = 0.;
#pragma omp parallel num_threads(nthreads)
{
#pragma omp for schedule(static)
for (size_t i=0; i<nx_dirty; ++i)
{
double fx = (i-nx_dirty/2.)*Pixsize_x();
fx *= fx;
for (size_t j=0; j<ny_dirty; ++j)
{
size_t i2 = nu-nx_dirty/2+i;
if (i2>=nu) i2-=nu;
size_t j2 = nv-ny_dirty/2+j;
if (j2>=nv) j2-=nv;
double fy = (j-ny_dirty/2.)*Pixsize_y();
auto ws = wscreen(fx, fy*fy, w, false);
ptmp[nv*i2+j2] = std::real(pdirty[ny_dirty*i + j])*ws*cfu[i]*cfv[j];
}
}
}
pocketfft::c2c({nu,nv}, strides, strides, {0,1}, pocketfft::FORWARD,
ptmp, ptmp, T(1), nthreads);
return complex<T>(cos(phase)*xn, sin(phase)*xn);
}
return tmp;
}
inline void getpix(T u_in, T v_in, T &u, T &v, int &iu0, int &iv0) const
{
u=fmodulo(u_in*psx, T(1))*nu,
......@@ -795,6 +724,40 @@ template<typename T> class GridderConfig
iv0 = int(v-w*0.5 + 1 + nv) - nv;
if (iv0+w>nv+nsafe) iv0 = nv+nsafe-w;
}
pyarr_c<complex<T>> apply_wscreen(const pyarr_c<complex<T>> &dirty_, double w, bool adjoint) const
{
checkArray(dirty_, "dirty", {nx_dirty, ny_dirty});
auto dirty = dirty_.data();
auto res_ = makeArray<complex<T>>({nx_dirty, ny_dirty});
auto res = res_.mutable_data();
double x0 = -0.5*nx_dirty*psx,
y0 = -0.5*ny_dirty*psy;
{
py::gil_scoped_release release;
#pragma omp parallel num_threads(nthreads)
{
#pragma omp for schedule(static)
for (size_t i=0; i<=nx_dirty/2; ++i)
{
double fx = x0+i*psx;
fx *= fx;
for (size_t j=0; j<=ny_dirty/2; ++j)
{
double fy = y0+j*psy;
auto ws = wscreen(fx, fy*fy, w, adjoint);
res[ny_dirty*i+j] = dirty[ny_dirty*i+j]*ws;
if (i>0)
res[ny_dirty*(nx_dirty-i)+j] = dirty[ny_dirty*(nx_dirty-i)+j]*ws;
if (j>0)
res[ny_dirty*i+(ny_dirty-j)] = dirty[ny_dirty*i+(ny_dirty-j)]*ws;
if ((i>0)&&(j>0))
res[ny_dirty*(nx_dirty-i)+(ny_dirty-j)] = dirty[ny_dirty*(nx_dirty-i)+(ny_dirty-j)]*ws;
}
}
}
}
return res_;
}
};
template<typename T, typename T2=complex<T>> class Helper
......@@ -1560,11 +1523,10 @@ PYBIND11_MODULE(nifty_gridder, m)
.def("grid2dirty", &GridderConfig<double>::grid2dirty,
grid2dirty_DS, "grid"_a)
.def("grid2dirty_c", &GridderConfig<double>::grid2dirty_c, "grid"_a)
.def("grid2dirty_wstacking", &GridderConfig<double>::grid2dirty_wstacking, "grid"_a, "w"_a)
.def("dirty2grid_wstacking", &GridderConfig<double>::dirty2grid_wstacking, "grid"_a, "w"_a)
.def("dirty2grid", &GridderConfig<double>::dirty2grid,
dirty2grid_DS, "dirty"_a)
.def("dirty2grid_c", &GridderConfig<double>::dirty2grid_c, "dirty"_a)
.def("apply_wscreen", &GridderConfig<double>::apply_wscreen, "dirty"_a, "w"_a, "adjoint"_a)
// pickle support
.def(py::pickle(
......
......@@ -28,15 +28,14 @@ if __name__ == '__main__':
w = 1000.2
ny, dy = nx, dx
conf = ng.GridderConfig(nx, ny, 1e-7, dx, dy)
x, y = conf.Nu(), conf.Nv()
x, y = conf.Nxdirty(), conf.Nydirty()
fld = np.random.randn(x, y) + 1j*np.random.randn(x, y)
func0 = lambda x: (_wscreen(nx, dy, w).conjugate()*conf.grid2dirty_c(x)
).real
func1 = lambda x: conf.grid2dirty_wstacking(x, w)
func0 = lambda x: _wscreen(nx, dy, w).conjugate()*x
func1 = lambda x: conf.apply_wscreen(x, w, True)
print(time_op(func0, fld), time_op(func1, fld))
fld = np.random.randn(nx, ny)
func0 = lambda x: conf.dirty2grid_c(_wscreen(nx, dx, w)*x.real)
func1 = lambda x: conf.dirty2grid_wstacking(x, w)
func0 = lambda x: _wscreen(nx, dx, w)*x.real
func1 = lambda x: conf.apply_wscreen(x, w, False)
print(time_op(func0, fld), time_op(func1, fld))
......@@ -173,10 +173,10 @@ def test_wstacking(nx, dx, w):
x, y = conf.Nu(), conf.Nv()
fld = np.random.randn(x, y) + 1j*np.random.randn(x, y)
res0 = (_wscreen(nx, dy, w).conjugate()*conf.grid2dirty_c(fld)).real
res1 = conf.grid2dirty_wstacking(fld, w)
res1 = conf.apply_wscreen(conf.grid2dirty_c(fld), w, True).real
np.testing.assert_allclose(res0, res1)
fld = np.random.randn(nx, ny)
res0 = conf.dirty2grid_c(_wscreen(nx, dx, w)*fld.real)
res1 = conf.dirty2grid_wstacking(fld, w)
res1 = conf.dirty2grid_c(conf.apply_wscreen(fld, w, False))
np.testing.assert_allclose(res0, res1)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment