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

test1

parent e3c40bc8
...@@ -659,41 +659,6 @@ template<typename T> class GridderConfig ...@@ -659,41 +659,6 @@ template<typename T> class GridderConfig
return res; 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 pyarr_c<T> dirty2grid(const pyarr_c<T> &dirty) const
{ {
checkArray(dirty, "dirty", {nx_dirty, ny_dirty}); checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
...@@ -742,50 +707,14 @@ template<typename T> class GridderConfig ...@@ -742,50 +707,14 @@ template<typename T> class GridderConfig
} }
return tmp; 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 double pi = 3.141592653589793238462643383279502884197;
constexpr complex<double> imagunit {0,1}; double n = cos(sqrt(x+y)), xn = 1./n;
double n = cos(sqrt(x+y)); double phase = 2*pi*w*(n-1);
complex<T> phase = 2*pi*w*imagunit*(n-1);
if (adjoint) phase *= -1; if (adjoint) phase *= -1;
return exp(phase)/n; return complex<T>(cos(phase)*xn, sin(phase)*xn);
}
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 tmp;
}
inline void getpix(T u_in, T v_in, T &u, T &v, int &iu0, int &iv0) const 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, u=fmodulo(u_in*psx, T(1))*nu,
...@@ -795,6 +724,40 @@ template<typename T> class GridderConfig ...@@ -795,6 +724,40 @@ template<typename T> class GridderConfig
iv0 = int(v-w*0.5 + 1 + nv) - nv; iv0 = int(v-w*0.5 + 1 + nv) - nv;
if (iv0+w>nv+nsafe) iv0 = nv+nsafe-w; 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 template<typename T, typename T2=complex<T>> class Helper
...@@ -1560,11 +1523,10 @@ PYBIND11_MODULE(nifty_gridder, m) ...@@ -1560,11 +1523,10 @@ PYBIND11_MODULE(nifty_gridder, m)
.def("grid2dirty", &GridderConfig<double>::grid2dirty, .def("grid2dirty", &GridderConfig<double>::grid2dirty,
grid2dirty_DS, "grid"_a) grid2dirty_DS, "grid"_a)
.def("grid2dirty_c", &GridderConfig<double>::grid2dirty_c, "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, .def("dirty2grid", &GridderConfig<double>::dirty2grid,
dirty2grid_DS, "dirty"_a) dirty2grid_DS, "dirty"_a)
.def("dirty2grid_c", &GridderConfig<double>::dirty2grid_c, "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 // pickle support
.def(py::pickle( .def(py::pickle(
......
...@@ -28,15 +28,14 @@ if __name__ == '__main__': ...@@ -28,15 +28,14 @@ if __name__ == '__main__':
w = 1000.2 w = 1000.2
ny, dy = nx, dx ny, dy = nx, dx
conf = ng.GridderConfig(nx, ny, 1e-7, dx, dy) 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) fld = np.random.randn(x, y) + 1j*np.random.randn(x, y)
func0 = lambda x: (_wscreen(nx, dy, w).conjugate()*conf.grid2dirty_c(x) func0 = lambda x: _wscreen(nx, dy, w).conjugate()*x
).real func1 = lambda x: conf.apply_wscreen(x, w, True)
func1 = lambda x: conf.grid2dirty_wstacking(x, w)
print(time_op(func0, fld), time_op(func1, fld)) print(time_op(func0, fld), time_op(func1, fld))
fld = np.random.randn(nx, ny) fld = np.random.randn(nx, ny)
func0 = lambda x: conf.dirty2grid_c(_wscreen(nx, dx, w)*x.real) func0 = lambda x: _wscreen(nx, dx, w)*x.real
func1 = lambda x: conf.dirty2grid_wstacking(x, w) func1 = lambda x: conf.apply_wscreen(x, w, False)
print(time_op(func0, fld), time_op(func1, fld)) print(time_op(func0, fld), time_op(func1, fld))
...@@ -173,10 +173,10 @@ def test_wstacking(nx, dx, w): ...@@ -173,10 +173,10 @@ def test_wstacking(nx, dx, w):
x, y = conf.Nu(), conf.Nv() x, y = conf.Nu(), conf.Nv()
fld = np.random.randn(x, y) + 1j*np.random.randn(x, y) fld = np.random.randn(x, y) + 1j*np.random.randn(x, y)
res0 = (_wscreen(nx, dy, w).conjugate()*conf.grid2dirty_c(fld)).real 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) np.testing.assert_allclose(res0, res1)
fld = np.random.randn(nx, ny) fld = np.random.randn(nx, ny)
res0 = conf.dirty2grid_c(_wscreen(nx, dx, w)*fld.real) 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) np.testing.assert_allclose(res0, res1)
Supports Markdown
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