Skip to content

GitLab

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

merge

parent 2fe1cc19
Hide whitespace changes
Inline Side-by-side
nifty_gridder.cc
View file @ c5fc5971
......@@ -131,7 +131,7 @@ template<typename T>
template<typename T>
using pyarr_c = py::array_t<T, py::array::c_style | py::array::forcecast>;
template<typename T> pyarr_c<T> makearray(const vector<size_t> &shape)
template<typename T> pyarr_c<T> makeArray(const vector<size_t> &shape)
{ return pyarr_c<T>(shape); }
size_t get_w(double epsilon)
......@@ -147,14 +147,49 @@ size_t get_w(double epsilon)
throw runtime_error("requested epsilon too small - minimum is 2e-13");
}
void checkArray(const py::array &arr, const char *aname,
const vector<size_t> &shape)
{
if (size_t(arr.ndim())!=shape.size())
{
cerr << "Array '" << aname << "' has " << arr.ndim() << " dimensions; "
"expected " << shape.size() << endl;
throw runtime_error("bad dimensionality");
}
for (size_t i=0; i<shape.size(); ++i)
if ((shape[i]!=0) && (size_t(arr.shape(i))!=shape[i]))
{
cerr << "Dimension " << i << " of array '" << aname << "' has size "
<< arr.shape(i) << "; expected " << shape[i] << endl;
throw runtime_error("bad array size");
}
}
template<typename T> pyarr_c<T> provideArray(py::object &in,
const vector<size_t> &shape)
{
if (in.is(py::none()))
{
auto tmp_ = makeArray<T>(shape);
size_t sz = size_t(tmp_.size());
auto tmp = tmp_.mutable_data();
for (size_t i=0; i<sz; ++i)
tmp[i] = T(0);
return tmp_;
}
auto tmp_ = in.cast<pyarr_c<T>>();
checkArray(tmp_, "temporary", shape);
return tmp_;
}
template<typename T> pyarr_c<T> complex2hartley
(const pyarr_c<complex<T>> &grid_)
{
myassert(grid_.ndim()==2, "grid array must be 2D");
checkArray(grid_, "grid", {0,0});
size_t nu = size_t(grid_.shape(0)), nv = size_t(grid_.shape(1));
auto grid = grid_.data();
auto res = makearray<T>({nu,nv});
auto res = makeArray<T>({nu,nv});
auto grid2 = res.mutable_data();
#pragma omp parallel for
for (size_t u=0; u<nu; ++u)
......@@ -175,11 +210,11 @@ template<typename T> pyarr_c<T> complex2hartley
template<typename T> pyarr_c<complex<T>> hartley2complex
(const pyarr_c<T> &grid_)
{
myassert(grid_.ndim()==2, "grid array must be 2D");
checkArray(grid_, "grid", {0, 0});
size_t nu = size_t(grid_.shape(0)), nv = size_t(grid_.shape(1));
auto grid = grid_.data();
auto res=makearray<complex<T>>({nu, nv});
auto res=makeArray<complex<T>>({nu, nv});
auto grid2 = res.mutable_data();
#pragma omp parallel for
for (size_t u=0; u<nu; ++u)
......@@ -265,9 +300,8 @@ template<typename T> class Baselines
public:
Baselines(const pyarr_c<T> &coord_, const pyarr_c<T> &lambda_)
{
myassert(coord_.ndim()==2, "coord array must be 2D");
myassert(coord_.shape(1)==3, "coord.shape[1] must be 3");
myassert(lambda_.ndim()==1, "lambda array must be 1D");
checkArray(coord_, "coord", {0, 3});
checkArray(lambda_, "lambda", {0});
nrows = coord_.shape(0);
nchan = lambda_.shape(0);
xlambda.resize(nchan);
......@@ -294,15 +328,13 @@ template<typename T> class Baselines
template<typename T2> pyarr_c<T2> ms2vis(const pyarr_c<T2> &ms_,
const pyarr_c<uint32_t> &idx_) const
{
myassert(idx_.ndim()==1, "idx array must be 1D");
myassert(ms_.ndim()==2, "ms array must be 2D");
myassert(size_t(ms_.shape(0))==nrows, "baselines 1st dim mismatch");
myassert(size_t(ms_.shape(1))==nchan, "baselines 2nd dim mismatch");
checkArray(idx_, "idx", {0});
checkArray(ms_, "ms", {nrows, nchan});
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.data();
auto ms = ms_.data();
auto res=makearray<T2>({nvis});
auto res=makeArray<T2>({nvis});
auto vis = res.mutable_data();
#pragma omp parallel for
for (size_t i=0; i<nvis; ++i)
......@@ -311,44 +343,21 @@ template<typename T> class Baselines
}
template<typename T2> pyarr_c<T2> vis2ms(const pyarr_c<T2> &vis_,
const pyarr_c<uint32_t> &idx_) const
const pyarr_c<uint32_t> &idx_, py::object &ms_in) const
{
myassert(idx_.ndim()==1, "idx array must be 1D");
myassert(vis_.ndim()==1, "vis array must be 1D");
checkArray(vis_, "vis", {0});
size_t nvis = size_t(vis_.shape(0));
myassert(int(nvis)==idx_.shape(0), "idx/vis size mismatch");
checkArray(idx_, "idx", {nvis});
auto idx = idx_.data();
auto vis = vis_.data();
auto res = makearray<T2>({nrows, nchan});
auto res = provideArray<T2>(ms_in, {nrows, nchan});
auto ms = res.mutable_data();
for (size_t i=0; i<nrows*nchan; ++i)
ms[i] = T2(0);
#pragma omp parallel for
for (size_t i=0; i<nvis; ++i)
ms[idx[i]] = vis[i];
return res;
}
template<typename T2> pyarr_c<T2> add_vis_to_ms(const pyarr_c<T2> &vis_,
const pyarr_c<uint32_t> &idx_, pyarr_c<T2> &ms_) const
{
myassert(idx_.ndim()==1, "idx array must be 1D");
myassert(vis_.ndim()==1, "vis array must be 1D");
myassert(ms_.ndim()==2, "ms array must be 2D");
myassert(size_t(ms_.shape(0))==nrows, "ms: bad 1st dimension");
myassert(size_t(ms_.shape(1))==nchan, "ms: bad 2nd dimension");
size_t nvis = size_t(vis_.shape(0));
myassert(int(nvis)==idx_.shape(0), "idx/vis size mismatch");
auto idx = idx_.data();
auto vis = vis_.data();
auto ms = ms_.mutable_data();
#pragma omp parallel for
for (size_t i=0; i<nvis; ++i)
ms[idx[i]] += vis[i];
return ms_;
}
};
constexpr int logsquare=4;
......@@ -396,13 +405,11 @@ template<typename T> class GridderConfig
T Beta() const { return beta; }
pyarr_c<T> grid2dirty(const pyarr_c<T> &grid) const
{
myassert(grid.ndim()==2, "grid must be a 2D array");
myassert(size_t(grid.shape(0))==nu, "bad 1st dimension");
myassert(size_t(grid.shape(1))==nv, "bad 2nd dimension");
auto tmp = makearray<T>({nu, nv});
checkArray(grid, "grid", {nu, nv});
auto tmp = makeArray<T>({nu, nv});
auto ptmp = tmp.mutable_data();
hartley2_2D<T>(grid, tmp);
auto res = makearray<T>({nx_dirty, ny_dirty});
auto res = makeArray<T>({nx_dirty, ny_dirty});
auto pout = res.mutable_data();
for (size_t i=0; i<nx_dirty; ++i)
for (size_t j=0; j<ny_dirty; ++j)
......@@ -417,15 +424,13 @@ template<typename T> class GridderConfig
}
pyarr_c<complex<T>> grid2dirty_c(const pyarr_c<complex<T>> &grid) const
{
myassert(grid.ndim()==2, "grid must be a 2D array");
myassert(size_t(grid.shape(0))==nu, "bad 1st dimension");
myassert(size_t(grid.shape(1))==nv, "bad 2nd dimension");
auto tmp = makearray<complex<T>>({nu, nv});
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), 0);
auto res = makearray<complex<T>>({nx_dirty, ny_dirty});
auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
auto pout = res.mutable_data();
for (size_t i=0; i<nx_dirty; ++i)
for (size_t j=0; j<ny_dirty; ++j)
......@@ -440,11 +445,9 @@ template<typename T> class GridderConfig
}
pyarr_c<T> dirty2grid(const pyarr_c<T> &dirty) const
{
myassert(dirty.ndim()==2, "dirty must be a 2D array");
myassert(size_t(dirty.shape(0))==nx_dirty, "bad 1st dimension");
myassert(size_t(dirty.shape(1))==ny_dirty, "bad 2nd dimension");
checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
auto pdirty = dirty.data();
auto tmp = makearray<T>({nu, nv});
auto tmp = makeArray<T>({nu, nv});
auto ptmp = tmp.mutable_data();
for (size_t i=0; i<nu*nv; ++i)
ptmp[i] = 0.;
......@@ -462,11 +465,9 @@ template<typename T> class GridderConfig
}
pyarr_c<complex<T>> dirty2grid_c(const pyarr_c<complex<T>> &dirty) const
{
myassert(dirty.ndim()==2, "dirty must be a 2D array");
myassert(size_t(dirty.shape(0))==nx_dirty, "bad 1st dimension");
myassert(size_t(dirty.shape(1))==ny_dirty, "bad 2nd dimension");
checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
auto pdirty = dirty.data();
auto tmp = makearray<complex<T>>({nu, nv});
auto tmp = makeArray<complex<T>>({nu, nv});
auto ptmp = tmp.mutable_data();
for (size_t i=0; i<nu*nv; ++i)
ptmp[i] = 0.;
......@@ -516,7 +517,7 @@ template<typename T> class Helper
{
if (bu0<-nsafe) return; // nothing written into buffer yet
#pragma omp critical
#pragma omp critical (gridder_writing_to_grid)
{
int idxu = (bu0+nu)%nu;
int idxv0 = (bv0+nv)%nv;
......@@ -594,15 +595,14 @@ template<typename T> pyarr_c<complex<T>> vis2grid_c(const Baselines<T> &baseline
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
const pyarr_c<complex<T>> &vis_)
{
myassert(idx_.ndim()==1, "idx array must be 1D");
myassert(vis_.ndim()==1, "vis must be 1D");
checkArray(vis_, "vis", {0});
size_t nvis = size_t(vis_.shape(0));
checkArray(idx_, "idx", {nvis});
auto vis=vis_.data();
myassert(vis_.shape(0)==idx_.shape(0), "bad vis dimension");
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.data();
size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = makearray<complex<T>>({nu, nv});
auto res = makeArray<complex<T>>({nu, nv});
auto grid = res.mutable_data();
for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
T beta = gconf.Beta();
......@@ -640,20 +640,69 @@ template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
const pyarr_c<complex<T>> &vis_)
{ return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_)); }
template<typename T> pyarr_c<complex<T>> ms2grid_c(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
const pyarr_c<complex<T>> &ms_)
{
auto nrows = baselines.Nrows();
auto nchan = baselines.Nchannels();
checkArray(ms_, "ms", {nrows, nchan});
checkArray(idx_, "idx", {0});
size_t nvis = size_t(idx_.shape(0));
auto ms = ms_.data();
auto idx = idx_.data();
size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = makeArray<complex<T>>({nu, nv});
auto grid = res.mutable_data();
for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
T beta = gconf.Beta();
size_t w = gconf.W();
#pragma omp parallel
{
Helper<T> hlp(gconf, grid, true);
T emb = exp(-2*beta);
const T * RESTRICT ku = hlp.kernel.data();
const T * RESTRICT kv = hlp.kernel.data()+w;
// Loop over sampling points
#pragma omp for schedule(guided,100)
for (size_t ipart=0; ipart<nvis; ++ipart)
{
UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
hlp.prep(coord.u, coord.v);
auto * RESTRICT ptr = hlp.p0;
auto v(ms[idx[ipart]]*emb);
for (size_t cu=0; cu<w; ++cu)
{
complex<T> tmp(v*ku[cu]);
for (size_t cv=0; cv<w; ++cv)
ptr[cv] += tmp*kv[cv];
ptr+=hlp.sv;
}
}
} // end of parallel region
return res;
}
template<typename T> pyarr_c<T> ms2grid(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
const pyarr_c<complex<T>> &ms_)
{ return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_)); }
template<typename T> pyarr_c<complex<T>> grid2vis_c(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
const pyarr_c<complex<T>> &grid_)
{
size_t nu=gconf.Nu(), nv=gconf.Nv();
myassert(idx_.ndim()==1, "idx array must be 1D");
checkArray(idx_, "idx", {0});
auto grid = grid_.data();
myassert(grid_.ndim()==2, "data must be 2D");
myassert(grid_.shape(0)==int(nu), "bad grid dimension");
myassert(grid_.shape(1)==int(nv), "bad grid dimension");
checkArray(grid_, "grid", {nu, nv});
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.data();
auto res = makearray<complex<T>>({nvis});
auto res = makeArray<complex<T>>({nvis});
auto vis = res.mutable_data();
T beta = gconf.Beta();
size_t w = gconf.W();
......@@ -692,6 +741,57 @@ template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
const pyarr_c<T> &grid_)
{ return grid2vis_c(baselines, gconf, idx_, hartley2complex(grid_)); }
template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
const pyarr_c<complex<T>> &grid_, py::object &ms_in)
{
size_t nu=gconf.Nu(), nv=gconf.Nv();
checkArray(idx_, "idx", {0});
auto grid = grid_.data();
checkArray(grid_, "grid", {nu, nv});
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.data();
auto res = provideArray<complex<T>>(ms_in,
{baselines.Nrows(), baselines.Nchannels()});
auto ms = res.mutable_data();
T beta = gconf.Beta();
size_t w = gconf.W();
// Loop over sampling points
#pragma omp parallel
{
Helper<T> hlp(gconf, grid, false);
T emb = exp(-2*beta);
const T * RESTRICT ku = hlp.kernel.data();
const T * RESTRICT kv = hlp.kernel.data()+w;
#pragma omp for schedule(guided,100)
for (size_t ipart=0; ipart<nvis; ++ipart)
{
UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
hlp.prep(coord.u, coord.v);
complex<T> r = 0;
auto * RESTRICT ptr = hlp.p0;
for (size_t cu=0; cu<w; ++cu)
{
complex<T> tmp(0);
for (size_t cv=0; cv<w; ++cv)
tmp += ptr[cv] * kv[cv];
r += tmp*ku[cu];
ptr += hlp.sv;
}
ms[idx[ipart]] += r*emb;
}
}
return res;
}
template<typename T> pyarr_c<complex<T>> grid2ms(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
const pyarr_c<T> &grid_, py::object &ms_in)
{ return grid2ms_c(baselines, gconf, idx_, hartley2complex(grid_), ms_in); }
template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<bool> &flags_, int chbegin,
int chend, T wmin, T wmax)
......@@ -701,9 +801,7 @@ template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
nsafe=gconf.Nsafe();
if (chbegin<0) chbegin=0;
if (chend<0) chend=nchan;
myassert(flags_.ndim()==2, "flags must be 2D");
myassert(size_t(flags_.shape(0))==nrow, "bad flags dimension");
myassert(size_t(flags_.shape(1))==nchan, "bad flags dimension");
checkArray(flags_, "flags", {nrow, nchan});
auto flags = flags_.data();
constexpr int side=1<<logsquare;
size_t nbu = (gconf.Nu()+1+side-1) >> logsquare,
......@@ -729,7 +827,7 @@ template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
for (size_t i=1; i<acc.size(); ++i)
acc[i] += acc[i-1];
auto res = makearray<uint32_t>({acc.back()});
auto res = makeArray<uint32_t>({acc.back()});
auto iout = res.mutable_data();
for (size_t irow=0, idx=0; irow<nrow; ++irow)
for (int ichan=chbegin; ichan<chend; ++ichan)
......@@ -785,24 +883,6 @@ np.array((nrows, nchannels), dtype=np.complex)
the measurement set's visibility data (0 where not covered by idx)
)""";
const char *BL_add_vis_to_ms_DS = R"""(
Adds a set of visibilities to an existing MS.
Parameters
==========
vis: np.array((nvis,), dtype=np.complex)
The visibility data for the index array
idx: np.array((nvis,), dtype=np.uint32)
the indices to be inserted
ms: np.array((nrows, nchannels), dtype=np.complex)
a MS
Returns
=======
np.array((nrows, nchannels), dtype=np.complex)
the updated MS
)""";
const char *GridderConfig_DS = R"""(
Class storing information related to the gridding/degridding process.
......@@ -932,10 +1012,7 @@ PYBIND11_MODULE(nifty_gridder, m)
.def ("Nrows",&Baselines<double>::Nrows)
.def ("Nchannels",&Baselines<double>::Nchannels)
.def ("ms2vis",&Baselines<double>::ms2vis<complex<double>>, BL_ms2vis_DS, "ms"_a, "idx"_a)
// .def ("ms2vis_f32",&Baselines<double>::ms2vis<float>, "ms"_a, "idx"_a)
.def ("vis2ms",&Baselines<double>::vis2ms<complex<double>>, BL_vis2ms_DS, "vis"_a, "idx"_a)
.def ("add_vis_to_ms",&Baselines<double>::add_vis_to_ms<complex<double>>, BL_add_vis_to_ms_DS,
"vis"_a, "idx"_a, "ms"_a.noconvert());
.def ("vis2ms",&Baselines<double>::vis2ms<complex<double>>, BL_vis2ms_DS, "vis"_a, "idx"_a, "ms_in"_a=py::none());
py::class_<GridderConfig<double>> (m, "GridderConfig", GridderConfig_DS)
.def(py::init<size_t, size_t, double, double, double>(),"nxdirty"_a,
"nydirty"_a, "epsilon"_a, "fov_x"_a, "fov_y"_a)
......@@ -948,30 +1025,11 @@ PYBIND11_MODULE(nifty_gridder, m)
m.def("getIndices", getIndices<double>, getIndices_DS, "baselines"_a, "gconf"_a,
"flags"_a, "chbegin"_a=-1, "chend"_a=-1, "wmin"_a=-1e30, "wmax"_a=1e30);
m.def("vis2grid",&vis2grid<double>, vis2grid_DS, "baselines"_a, "gconf"_a, "idx"_a, "vis"_a);
m.def("ms2grid",&ms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a);
m.def("grid2vis",&grid2vis<double>, grid2vis_DS, "baselines"_a, "gconf"_a, "idx"_a, "grid"_a);
m.def("grid2ms",&grid2ms<double>, "baselines"_a, "gconf"_a, "idx"_a, "grid"_a, "ms_in"_a=py::none());
m.def("vis2grid_c",&vis2grid_c<double>, "baselines"_a, "gconf"_a, "idx"_a, "vis"_a);
m.def("ms2grid_c",&ms2grid_c<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a);
m.def("grid2vis_c",&grid2vis_c<double>, "baselines"_a, "gconf"_a, "idx"_a, "grid"_a);
#if 0
py::class_<Baselines<float>> (m, "Baselines_f", Baselines_DS)
.def(py::init<const pyarr_c<float> &, const pyarr_c<float> &>(),
"coord"_a, "lambda"_a)
.def ("Nrows",&Baselines<float>::Nrows)
.def ("Nchannels",&Baselines<float>::Nchannels)
.def ("ms2vis",&Baselines<float>::ms2vis<complex<float>>, BL_ms2vis_DS, "ms"_a, "idx"_a)
.def ("vis2ms",&Baselines<float>::vis2ms<complex<float>>, BL_vis2ms_DS, "vis"_a, "idx"_a)
.def ("add_vis_to_ms",&Baselines<float>::add_vis_to_ms<complex<float>>, BL_add_vis_to_ms_DS,
"vis"_a, "idx"_a, "ms"_a.noconvert());
py::class_<GridderConfig<float>> (m, "GridderConfig_f")
.def(py::init<size_t, size_t, float, float, float>(),"nxdirty"_a,
"nydirty"_a, "epsilon"_a, "fov_x"_a, "fov_y"_a)
.def("Nu", &GridderConfig<float>::Nu)
.def("Nv", &GridderConfig<float>::Nv)
.def("grid2dirty", &GridderConfig<float>::grid2dirty, "grid"_a)
.def("dirty2grid", &GridderConfig<float>::dirty2grid, "dirty"_a);
m.def("getIndices_f", getIndices<float>, "baselines"_a, "gconf"_a,
"flags"_a, "chbegin"_a=-1, "chend"_a=-1, "wmin"_a=-1e30, "wmax"_a=1e30);
m.def("vis2grid_f",&vis2grid<float>, "baselines"_a, "gconf"_a, "idx"_a, "vis"_a);
m.def("grid2vis_f",&grid2vis<float>, "baselines"_a, "gconf"_a, "idx"_a, "grid"_a);
#endif
m.def("grid2ms_c",&grid2ms_c<double>, "baselines"_a, "gconf"_a, "idx"_a, "grid"_a, "ms_in"_a=py::none());
}
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