Skip to content

GitLab

Commit 15ade055 authored by Martin Reinecke's avatar Martin Reinecke
Browse files

reintroduce real-only functions

parent 9728e8e8
Hide whitespace changes
Inline Side-by-side
gridder_cxx.h
View file @ 15ade055
......@@ -85,9 +85,7 @@ template<typename T, size_t ndim> class mav
}
operator mav<const T, ndim>() const
{ return mav<const T, ndim>(d,shp,str); }
T &operator[](size_t i)
{ return operator()(i); }
const T &operator[](size_t i) const
T &operator[](size_t i) const
{ return operator()(i); }
T &operator()(size_t i) const
{
......@@ -108,9 +106,7 @@ template<typename T, size_t ndim> class mav
return res;
}
ptrdiff_t stride(size_t i) const { return str[i]; }
T *data()
{ return d; }
const T *data() const
T *data() const
{ return d; }
bool contiguous() const
{
......@@ -269,6 +265,72 @@ struct RowChan
size_t row, chan;
};
template<typename T> void complex2hartley
(const const_mav<complex<T>, 2> &grid, const mav<T,2> &grid2, size_t nthreads)
{
myassert(grid.shape()==grid2.shape(), "shape mismatch");
size_t nu=grid.shape(0), nv=grid.shape(1);
#pragma omp parallel for num_threads(nthreads)
for (size_t u=0; u<nu; ++u)
{
size_t xu = (u==0) ? 0 : nu-u;
for (size_t v=0; v<nv; ++v)
{
size_t xv = (v==0) ? 0 : nv-v;
grid2(u,v) += T(0.5)*(grid( u, v).real()+grid( u, v).imag()+
grid(xu,xv).real()-grid(xu,xv).imag());
}
}
}
template<typename T> void hartley2complex
(const const_mav<T,2> &grid, const mav<complex<T>,2> &grid2, size_t nthreads)
{
myassert(grid.shape()==grid2.shape(), "shape mismatch");
size_t nu=grid.shape(0), nv=grid.shape(1);
#pragma omp parallel for num_threads(nthreads)
for (size_t u=0; u<nu; ++u)
{
size_t xu = (u==0) ? 0 : nu-u;
for (size_t v=0; v<nv; ++v)
{
size_t xv = (v==0) ? 0 : nv-v;
T v1 = T(0.5)*grid( u, v);
T v2 = T(0.5)*grid(xu,xv);
grid2(u,v) = std::complex<T>(v1+v2, v1-v2);
}
}
}
template<typename T> void hartley2_2D(const const_mav<T,2> &in, const mav<T,2> &out, size_t nthreads)
{
myassert(in.shape()==out.shape(), "shape mismatch");
size_t nu=in.shape(0), nv=in.shape(1);
ptrdiff_t sz=ptrdiff_t(sizeof(T));
pocketfft::stride_t stri{sz*in.stride(0), sz*in.stride(1)};
pocketfft::stride_t stro{sz*out.stride(0), sz*out.stride(1)};
auto d_i = in.data();
auto ptmp = out.data();
pocketfft::r2r_separable_hartley({nu, nv}, stri, stro, {0,1}, d_i, ptmp, T(1),
nthreads);
#pragma omp parallel for num_threads(nthreads)
for(size_t i=1; i<(nu+1)/2; ++i)
for(size_t j=1; j<(nv+1)/2; ++j)
{
T a = ptmp[i*nv+j];
T b = ptmp[(nu-i)*nv+j];
T c = ptmp[i*nv+nv-j];
T d = ptmp[(nu-i)*nv+nv-j];
ptmp[i*nv+j] = T(0.5)*(a+b+c-d);
ptmp[(nu-i)*nv+j] = T(0.5)*(a+b+d-c);
ptmp[i*nv+nv-j] = T(0.5)*(a+c+d-b);
ptmp[(nu-i)*nv+nv-j] = T(0.5)*(b+c+d-a);
}
}
template<typename T> class EC_Kernel
{
protected:
......@@ -499,12 +561,30 @@ template<typename T> class GridderConfig
img2(i,j) = img(i,j)*cfu[i]*cfv[j];
}
void grid2dirty(const const_mav<T,2> &grid, const mav<T,2> &dirty) const
{
checkShape(grid.shape(), {nu,nv});
checkShape(dirty.shape(), {nx_dirty,ny_dirty});
tmpStorage<T,2> tmpdat({nu,nv});
auto tmav = tmpdat.getMav();
hartley2_2D<T>(grid, tmav, nthreads);
for (size_t i=0; i<nx_dirty; ++i)
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;
dirty(i,j) = tmav(i2,j2)*cfu[i]*cfv[j];
}
}
void grid2dirty_c(const mav<const complex<T>,2> &grid, mav<complex<T>,2> &dirty) const
{
checkShape(grid.shape(), {nu,nv});
checkShape(dirty.shape(), {nx_dirty,ny_dirty});
vector<complex<T>> tmpdat(nu*nv);
auto tmp=mav<complex<T>,2>(tmpdat.data(),{nu,nv},{ptrdiff_t(nv),ptrdiff_t(1)});
tmpStorage<complex<T>,2> tmpdat({nu,nv});
auto tmp = tmpdat.getMav();
constexpr auto sc = ptrdiff_t(sizeof(complex<T>));
pocketfft::c2c({nu,nv},{grid.stride(0)*sc,grid.stride(1)*sc},
{tmp.stride(0)*sc, tmp.stride(1)*sc}, {0,1}, pocketfft::BACKWARD,
......@@ -520,14 +600,29 @@ template<typename T> class GridderConfig
}
}
void dirty2grid(const const_mav<T,2> &dirty, mav<T,2> &grid) const
{
checkShape(dirty.shape(), {nx_dirty, ny_dirty});
checkShape(grid.shape(), {nu, nv});
grid.fill(0);
for (size_t i=0; i<nx_dirty; ++i)
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;
grid(i2,j2) = dirty(i,j)*cfu[i]*cfv[j];
}
hartley2_2D<T>(grid, grid, nthreads);
}
void dirty2grid_c(const const_mav<complex<T>,2> &dirty,
mav<complex<T>,2> &grid) const
{
checkShape(dirty.shape(), {nx_dirty, ny_dirty});
checkShape(grid.shape(), {nu, nv});
for (size_t i=0; i<nu; ++i)
for (size_t j=0; j<nv; ++j)
grid(i,j)=0.;
grid.fill(0);
for (size_t i=0; i<nx_dirty; ++i)
for (size_t j=0; j<ny_dirty; ++j)
{
......
nifty_gridder.cc
View file @ 15ade055
......@@ -233,6 +233,34 @@ template<typename T> class PyBaselines: public Baselines<T>
}
};
constexpr auto grid2dirty_DS = R"""(
Converts from UV grid to dirty image (FFT, cropping, correction)
Parameters
==========
grid: np.array((nu, nv), dtype=np.float64)
gridded UV data
Returns
=======
nd.array((nxdirty, nydirty), dtype=np.float64)
the dirty image
)""";
constexpr auto dirty2grid_DS = R"""(
Converts from a dirty image to a UV grid (correction, padding, FFT)
Parameters
==========
dirty: nd.array((nxdirty, nydirty), dtype=np.float64)
the dirty image
Returns
=======
np.array((nu, nv), dtype=np.float64)
gridded UV data
)""";
constexpr auto apply_taper_DS = R"""(
Applies the taper (or its inverse) to an image
......@@ -318,6 +346,17 @@ template<typename T> class PyGridderConfig: public GridderConfig<T>
}
return res;
}
pyarr<T> grid2dirty(const pyarr<T> &grid) const
{
auto res = makeArray<T>({nx_dirty, ny_dirty});
auto grid2=make_const_mav<2>(grid);
auto res2=make_mav<2>(res);
{
py::gil_scoped_release release;
GridderConfig<T>::grid2dirty(grid2,res2);
}
return res;
}
pyarr<complex<T>> grid2dirty_c(const pyarr<complex<T>> &grid) const
{
auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
......@@ -330,6 +369,17 @@ template<typename T> class PyGridderConfig: public GridderConfig<T>
return res;
}
pyarr<T> dirty2grid(const pyarr<T> &dirty) const
{
auto dirty2 = make_const_mav<2>(dirty);
auto grid = makeArray<T>({nu, nv});
auto grid2=make_mav<2>(grid);
{
py::gil_scoped_release release;
GridderConfig<T>::dirty2grid(dirty2, grid2);
}
return grid;
}
pyarr<complex<T>> dirty2grid_c(const pyarr<complex<T>> &dirty) const
{
auto dirty2 = make_const_mav<2>(dirty);
......@@ -422,6 +472,15 @@ Returns
np.array((nu,nv), dtype=np.float64):
the gridded visibilities (made real by making use of Hermitian symmetry)
)""";
template<typename T> pyarr<T> Pyvis2grid(const PyBaselines<T> &baselines,
const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr<complex<T>> &vis_, py::object &grid_in, const py::object &wgt_)
{
auto tmp=Pyvis2grid_c(baselines, gconf, idx_, vis_, None, wgt_);
auto grd=provideCArray<T>(grid_in,{gconf.Nu(), gconf.Nv()});
gridder::detail::complex2hartley(make_const_mav<2>(tmp), make_mav<2>(grd), gconf.Nthreads());
return grd;
}
constexpr auto ms2grid_c_DS = R"""(
Grids measurement set data onto a UV grid
......@@ -467,6 +526,18 @@ template<typename T> pyarr<complex<T>> Pyms2grid_c(
return res;
}
template<typename T> pyarr<T> Pyms2grid(
const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_,
py::object &grid_in, const py::object &wgt_)
{
auto tmp = Pyms2grid_c(baselines, gconf, idx_, ms_, None, wgt_);
auto res_ = provideCArray<T>(grid_in, {gconf.Nu(), gconf.Nv()});
auto res = make_mav<2>(res_);
gridder::detail::complex2hartley(make_const_mav<2>(tmp), res, gconf.Nthreads());
return res_;
}
template<typename T> pyarr<complex<T>> Pygrid2vis_c(
const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &grid_,
......@@ -508,6 +579,15 @@ Returns
np.array((nvis,), dtype=np.complex)
The degridded visibility data
)""";
template<typename T> pyarr<complex<T>> Pygrid2vis(const PyBaselines<T> &baselines,
const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr<T> &grid_, const py::object &wgt_)
{
auto tmp=makeArray<complex<T>>({gconf.Nu(), gconf.Nv()});
gridder::detail::hartley2complex(make_const_mav<2>(grid_),make_mav<2>(tmp), gconf.Nthreads());
return Pygrid2vis_c(baselines, gconf, idx_, tmp, wgt_);
}
template<typename T> pyarr<complex<T>> Pygrid2ms_c(
const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &grid_,
......@@ -528,6 +608,16 @@ template<typename T> pyarr<complex<T>> Pygrid2ms_c(
return res;
}
template<typename T> pyarr<complex<T>> Pygrid2ms(const PyBaselines<T> &baselines,
const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr<T> &grid_, py::object &ms_in, const py::object &wgt_)
{
auto grid2_ = makeArray<complex<T>>({gconf.Nu(), gconf.Nv()});
auto grid2 = make_mav<2>(grid2_);
gridder::detail::hartley2complex(make_const_mav<2>(grid_), grid2, gconf.Nthreads());
return Pygrid2ms_c(baselines, gconf, idx_, grid2_, ms_in, wgt_);
}
template<typename T> pyarr<complex<T>> Pyapply_holo(
const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &grid_,
......@@ -593,13 +683,13 @@ np.array((nvis,), dtype=np.uint32)
)""";
template<typename T> pyarr<uint32_t> PygetIndices(const PyBaselines<T> &baselines,
const PyGridderConfig<T> &gconf, const pyarr<bool> &flags_, int chbegin,
int chend, T wmin, T wmax, size_t nthreads)
int chend, T wmin, T wmax)
{
size_t nidx;
auto flags = make_const_mav<2>(flags_);
{
py::gil_scoped_release release;
nidx = getIdxSize(baselines, flags, chbegin, chend, wmin, wmax, nthreads);
nidx = getIdxSize(baselines, flags, chbegin, chend, wmin, wmax, gconf.Nthreads());
}
auto res = makeArray<uint32_t>({nidx});
auto res2 = make_mav<1>(res);
......@@ -734,7 +824,11 @@ PYBIND11_MODULE(nifty_gridder, m)
.def("Supp", &PyGridderConfig<double>::Supp)
.def("apply_taper", &PyGridderConfig<double>::apply_taper, apply_taper_DS,
"img"_a, "divide"_a=false)
.def("grid2dirty", &PyGridderConfig<double>::grid2dirty,
grid2dirty_DS, "grid"_a)
.def("grid2dirty_c", &PyGridderConfig<double>::grid2dirty_c, "grid"_a)
.def("dirty2grid", &PyGridderConfig<double>::dirty2grid,
dirty2grid_DS, "dirty"_a)
.def("dirty2grid_c", &PyGridderConfig<double>::dirty2grid_c, "dirty"_a)
.def("apply_wscreen", &PyGridderConfig<double>::apply_wscreen,
apply_wscreen_DS, "dirty"_a, "w"_a, "adjoint"_a)
......@@ -759,7 +853,15 @@ PYBIND11_MODULE(nifty_gridder, m)
}));
m.def("getIndices", PygetIndices<double>, getIndices_DS, "baselines"_a,
"gconf"_a, "flags"_a, "chbegin"_a=-1, "chend"_a=-1,
"wmin"_a=-1e30, "wmax"_a=1e30, "nthreads"_a=1);
"wmin"_a=-1e30, "wmax"_a=1e30);
m.def("vis2grid",&Pyvis2grid<double>, vis2grid_DS, "baselines"_a, "gconf"_a,
"idx"_a, "vis"_a, "grid_in"_a=None, "wgt"_a=None);
m.def("ms2grid",&Pyms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a,
"grid_in"_a=None, "wgt"_a=None);
m.def("grid2vis",&Pygrid2vis<double>, grid2vis_DS, "baselines"_a, "gconf"_a,
"idx"_a, "grid"_a, "wgt"_a=None);
m.def("grid2ms",&Pygrid2ms<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "ms_in"_a=None, "wgt"_a=None);
m.def("vis2grid_c",&Pyvis2grid_c<double>, vis2grid_c_DS, "baselines"_a,
"gconf"_a, "idx"_a, "vis"_a, "grid_in"_a=None, "wgt"_a=None);
m.def("ms2grid_c",&Pyms2grid_c<double>, ms2grid_c_DS, "baselines"_a, "gconf"_a,
......@@ -778,27 +880,27 @@ PYBIND11_MODULE(nifty_gridder, m)
"idx"_a, "dirty"_a);
m.def("full_gridding",&Pyfull_gridding<double>,"uvw"_a,"freq"_a,"ms"_a,
"wgt"_a,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
"wgt"_a=None,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
"nthreads"_a=1, "verbosity"_a=0);
m.def("full_gridding_f",&Pyfull_gridding<float>,"uvw"_a,"freq"_a,"ms"_a,
"wgt"_a,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
"wgt"_a=None,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
"nthreads"_a=1, "verbosity"_a=0);
m.def("full_degridding",&Pyfull_degridding<double>,"uvw"_a,"freq"_a,"dirty"_a,
"wgt"_a,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a,"nthreads"_a=1,
"wgt"_a=None,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a,"nthreads"_a=1,
"verbosity"_a=0);
m.def("full_degridding_f",&Pyfull_degridding<float>,"uvw"_a,"freq"_a,"dirty"_a,
"wgt"_a,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a, "nthreads"_a=1,
"wgt"_a=None,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a, "nthreads"_a=1,
"verbosity"_a=0);
m.def("gridding",&Pygridding<double>,"uvw"_a,"freq"_a,"ms"_a,
"wgt"_a,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
"wgt"_a=None,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
"nthreads"_a=1, "verbosity"_a=0);
m.def("gridding_f",&Pygridding<float>,"uvw"_a,"freq"_a,"ms"_a,
"wgt"_a,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
"wgt"_a=None,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
"nthreads"_a=1, "verbosity"_a=0);
m.def("degridding",&Pydegridding<double>,"uvw"_a,"freq"_a,"dirty"_a,
"wgt"_a,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a,"nthreads"_a=1,
"wgt"_a=None,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a,"nthreads"_a=1,
"verbosity"_a=0);
m.def("degridding_f",&Pydegridding<float>,"uvw"_a,"freq"_a,"dirty"_a,
"wgt"_a,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a, "nthreads"_a=1,
"wgt"_a=None,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a, "nthreads"_a=1,
"verbosity"_a=0);
}
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