Skip to content
GitLab
Commit 50378509 authored by Martin Reinecke's avatar Martin Reinecke
Browse files

first try

parent fc1b38de
Show whitespace changes
Inline Side-by-side
nifty_gridder.cc
View file @ 50378509
......@@ -195,7 +195,7 @@ void checkArray(const py::array &arr, const char *aname,
}
}
template<typename T> pyarr<T> provideArray(py::object &in,
template<typename T> pyarr<T> provideArray(const py::object &in,
const vector<size_t> &shape)
{
if (in.is_none())
......@@ -212,6 +212,16 @@ template<typename T> pyarr<T> provideArray(py::object &in,
return tmp_;
}
template<typename T> pyarr<T> providePotentialArray(const py::object &in,
const vector<size_t> &shape)
{
if (in.is_none())
return makeArray<T>(vector<size_t>(shape.size(), 0));
auto tmp_ = in.cast<pyarr<T>>();
checkArray(tmp_, "temporary", shape);
return tmp_;
}
template<typename T> pyarr_c<T> provideCArray(py::object &in,
const vector<size_t> &shape)
{
......@@ -914,6 +924,8 @@ vis: np.array((nvis,), dtype=np.complex)
The visibility data for the index array
grid_in: np.array((nu,nv), dtype=np.complex128), optional
If present, the result is added to this array.
wgt: np.array((nvis,), dtype=np.float64), optional
If present, visibilities are multiplied by the corresponding entries.
Returns
=======
......@@ -923,13 +935,16 @@ np.array((nu,nv), dtype=np.complex128):
template<typename T> pyarr_c<complex<T>> vis2grid_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &vis_,
py::object &grid_in)
py::object &grid_in, const py::object &wgt_)
{
checkArray(vis_, "vis", {0});
size_t nvis = size_t(vis_.shape(0));
checkArray(idx_, "idx", {nvis});
auto vis=vis_.template unchecked<1>();
auto idx = idx_.template unchecked<1>();
pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {nvis});
bool have_wgt = wgt2.size()>0;
auto wgt = wgt2.template unchecked<1>();
size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = provideCArray<complex<T>>(grid_in, {nu, nv});
......@@ -955,6 +970,8 @@ template<typename T> pyarr_c<complex<T>> vis2grid_c(
hlp.prep(coord.u, coord.v);
auto * RESTRICT ptr = hlp.p0w;
auto v(vis(ipart)*emb);
if (have_wgt)
v*=wgt(ipart);
for (size_t cu=0; cu<w; ++cu)
{
complex<T> tmp(v*ku[cu]);
......@@ -984,6 +1001,8 @@ vis: np.array((nvis,), dtype=np.complex)
The visibility data for the index array
grid_in: np.array((nu,nv), dtype=np.float64), optional
If present, the result is added to this array.
wgt: np.array((nvis,), dtype=np.float64), optional
If present, visibilities are multiplied by the corresponding entries.
Returns
=======
......@@ -992,8 +1011,8 @@ np.array((nu,nv), dtype=np.float64):
)""";
template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr<complex<T>> &vis_, py::object &grid_in)
{ return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_, None), grid_in); }
const pyarr<complex<T>> &vis_, py::object &grid_in, const py::object &wgt_)
{ return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_, None, wgt_), grid_in); }
constexpr auto ms2grid_c_DS = R"""(
Grids measurement set data onto a UV grid
......@@ -1011,6 +1030,8 @@ ms: np.array((nrows, nchannels), dtype=np.complex128)
the measurement set.
grid_in: np.array((nu,nv), dtype=np.complex128), optional
If present, the result is added to this array.
wgt: np.array((nrows, nchannels), dtype=np.float64), optional
If present, visibilities are multiplied by the corresponding entries.
Returns
=======
......@@ -1020,7 +1041,7 @@ np.array((nu,nv), dtype=np.complex128):
template<typename T> pyarr_c<complex<T>> ms2grid_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_,
py::object &grid_in)
py::object &grid_in, const py::object &wgt_)
{
auto nrows = baselines.Nrows();
auto nchan = baselines.Nchannels();
......@@ -1029,6 +1050,9 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c(
size_t nvis = size_t(idx_.shape(0));
auto ms = ms_.template unchecked<2>();
auto idx = idx_.template unchecked<1>();
auto wgt2 = providePotentialArray<T>(wgt_, {nrows, nchan});
bool have_wgt = wgt2.size()>0;
auto wgt = wgt2.template unchecked<2>();
size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = provideCArray<complex<T>>(grid_in, {nu, nv});
......@@ -1057,6 +1081,8 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c(
hlp.prep(coord.u, coord.v);
auto * RESTRICT ptr = hlp.p0w;
auto v(ms(row,chan)*emb);
if (have_wgt)
v*=wgt(row, chan);
for (size_t cu=0; cu<w; ++cu)
{
complex<T> tmp(v*ku[cu]);
......@@ -1072,73 +1098,13 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c(
template<typename T> pyarr_c<T> ms2grid(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr<complex<T>> &ms_, py::object &grid_in)
{ return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_, None), grid_in); }
template<typename T> pyarr_c<complex<T>> ms2grid_c_wgt(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_,
const pyarr<T> &wgt_, py::object &grid_in)
{
auto nrows = baselines.Nrows();
auto nchan = baselines.Nchannels();
checkArray(wgt_, "wgt", {nrows, nchan});
checkArray(ms_, "ms", {nrows, nchan});
checkArray(idx_, "idx", {0});
size_t nvis = size_t(idx_.shape(0));
auto ms = ms_.template unchecked<2>();
auto wgt = wgt_.template unchecked<2>();
auto idx = idx_.template unchecked<1>();
size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = provideArray<complex<T>>(grid_in, {nu, nv});
auto grid = res.mutable_data();
{
py::gil_scoped_release release;
T beta = gconf.Beta();
size_t w = gconf.W();
#pragma omp parallel num_threads(nthreads)
{
Helper<T> hlp(gconf, nullptr, grid);
T emb = exp(-2*beta);
int jump = hlp.lineJump();
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)
{
auto tidx = idx(ipart);
auto row = tidx/nchan;
auto chan = tidx-row*nchan;
UVW<T> coord = baselines.effectiveCoord(tidx);
hlp.prep(coord.u, coord.v);
auto * RESTRICT ptr = hlp.p0w;
auto v(ms(row,chan)*(emb*wgt(row, chan)));
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+=jump;
}
}
} // end of parallel region
}
return res;
}
template<typename T> pyarr_c<T> ms2grid_wgt(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr<complex<T>> &ms_, const pyarr<T> &wgt_,
py::object &grid_in)
{ return complex2hartley(ms2grid_c_wgt(baselines, gconf, idx_, ms_, wgt_, None), grid_in); }
const pyarr<complex<T>> &ms_, py::object &grid_in, const py::object &wgt_)
{ return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_, None, wgt_), grid_in); }
template<typename T> pyarr_c<complex<T>> grid2vis_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_)
const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_,
const py::object &wgt_)
{
size_t nu=gconf.Nu(), nv=gconf.Nv();
checkArray(idx_, "idx", {0});
......@@ -1146,6 +1112,9 @@ template<typename T> pyarr_c<complex<T>> grid2vis_c(
checkArray(grid_, "grid", {nu, nv});
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.template unchecked<1>();
auto wgt2 = providePotentialArray<T>(wgt_, {nvis});
bool have_wgt = wgt2.size()>0;
auto wgt = wgt2.template unchecked<1>();
auto res = makeArray<complex<T>>({nvis});
auto vis = res.mutable_data();
......@@ -1178,6 +1147,7 @@ template<typename T> pyarr_c<complex<T>> grid2vis_c(
r += tmp*ku[cu];
ptr += jump;
}
if (have_wgt) r*=wgt[ipart];
vis[ipart] = r*emb;
}
}
......@@ -1199,6 +1169,8 @@ idx: np.array((nvis,), dtype=np.uint32)
the indices for the entries to be degridded
grid: np.array((nu,nv), dtype=np.float64):
the gridded visibilities (made real by making use of Hermitian symmetry)
wgt: np.array((nvis,), dtype=np.float64), optional
If present, visibilities are multiplied by the corresponding entries.
Returns
=======
......@@ -1207,12 +1179,12 @@ np.array((nvis,), dtype=np.complex)
)""";
template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr_c<T> &grid_)
{ return grid2vis_c(baselines, gconf, idx_, hartley2complex(grid_)); }
const pyarr_c<T> &grid_, const py::object &wgt_)
{ return grid2vis_c(baselines, gconf, idx_, hartley2complex(grid_), wgt_); }
template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr_c<complex<T>> &grid_, py::object &ms_in)
const pyarr_c<complex<T>> &grid_, py::object &ms_in, const py::object &wgt_)
{
auto nrows = baselines.Nrows();
auto nchan = baselines.Nchannels();
......@@ -1222,6 +1194,9 @@ template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines
checkArray(grid_, "grid", {nu, nv});
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.template unchecked<1>();
auto wgt2 = providePotentialArray<T>(wgt_, {nrows, nchan});
bool have_wgt = wgt2.size()>0;
auto wgt = wgt2.template unchecked<2>();
auto res = provideArray<complex<T>>(ms_in, {nrows, nchan});
auto ms = res.template mutable_unchecked<2>();
......@@ -1257,6 +1232,9 @@ template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines
r += tmp*ku[cu];
ptr += jump;
}
r*=emb;
if (have_wgt)
r*=wgt(row, chan);
ms(row,chan) += r*emb;
}
}
......@@ -1266,78 +1244,13 @@ template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines
template<typename T> pyarr_c<complex<T>> grid2ms(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<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<complex<T>> grid2ms_c_wgt(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_,
const pyarr<T> &wgt_, py::object &ms_in)
{
auto nrows = baselines.Nrows();
auto nchan = baselines.Nchannels();
checkArray(wgt_, "wgt", {nrows, nchan});
auto wgt = wgt_.template unchecked<2>();
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_.template unchecked<1>();
auto res = provideArray<complex<T>>(ms_in, {nrows, nchan});
auto ms = res.template mutable_unchecked<2>();
{
py::gil_scoped_release release;
T beta = gconf.Beta();
size_t w = gconf.W();
// Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
Helper<T> hlp(gconf, grid, nullptr);
T emb = exp(-2*beta);
int jump = hlp.lineJump();
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)
{
auto tidx = idx(ipart);
auto row = tidx/nchan;
auto chan = tidx-row*nchan;
UVW<T> coord = baselines.effectiveCoord(tidx);
hlp.prep(coord.u, coord.v);
complex<T> r = 0;
const auto * RESTRICT ptr = hlp.p0r;
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 += jump;
}
ms(row,chan) += r*(emb*wgt(row,chan));
}
}
}
return res;
}
template<typename T> pyarr_c<complex<T>> grid2ms_wgt(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr_c<T> &grid_, const pyarr<T> &wgt_,
py::object &ms_in)
{
return grid2ms_c_wgt(baselines, gconf, idx_, hartley2complex(grid_), wgt_,
ms_in);
}
const pyarr_c<T> &grid_, py::object &ms_in, const py::object &wgt_)
{ return grid2ms_c(baselines, gconf, idx_, hartley2complex(grid_), ms_in, wgt_); }
template<typename T> pyarr_c<complex<T>> apply_holo(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_)
const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_,
const py::object &wgt_)
{
size_t nu=gconf.Nu(), nv=gconf.Nv();
checkArray(idx_, "idx", {0});
......@@ -1345,6 +1258,9 @@ template<typename T> pyarr_c<complex<T>> apply_holo(
checkArray(grid_, "grid", {nu, nv});
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.template unchecked<1>();
pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {nvis});
bool have_wgt = wgt2.size()>0;
auto wgt = wgt2.template unchecked<1>();
auto res = makeArray<complex<T>>({nu, nv});
auto ogrid = res.mutable_data();
......@@ -1365,7 +1281,8 @@ template<typename T> pyarr_c<complex<T>> apply_holo(
#pragma omp for schedule(guided,100)
for (size_t ipart=0; ipart<nvis; ++ipart)
{
UVW<T> coord = baselines.effectiveCoord(idx(ipart));
auto tidx = idx(ipart);
UVW<T> coord = baselines.effectiveCoord(tidx);
hlp.prep(coord.u, coord.v);
complex<T> r = 0;
const auto * RESTRICT ptr = hlp.p0r;
......@@ -1378,6 +1295,11 @@ template<typename T> pyarr_c<complex<T>> apply_holo(
ptr += jump;
}
r*=emb*emb;
if (have_wgt)
{
auto twgt = wgt(ipart);
r*=twgt*twgt;
}
auto * RESTRICT wptr = hlp.p0w;
for (size_t cu=0; cu<w; ++cu)
{
......@@ -1394,7 +1316,7 @@ template<typename T> pyarr_c<complex<T>> apply_holo(
template<typename T> pyarr_c<T> get_correlations(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, int du, int dv)
const pyarr<uint32_t> &idx_, int du, int dv, py::object &wgt_)
{
size_t nu=gconf.Nu(), nv=gconf.Nv();
size_t w = gconf.W();
......@@ -1403,6 +1325,9 @@ template<typename T> pyarr_c<T> get_correlations(
checkArray(idx_, "idx", {0});
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.template unchecked<1>();
pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {nvis});
bool have_wgt = wgt2.size()>0;
auto wgt = wgt2.template unchecked<1>();
auto res = makeArray<T>({nu, nv});
auto ogrid = res.mutable_data();
......@@ -1436,9 +1361,15 @@ template<typename T> pyarr_c<T> get_correlations(
UVW<T> coord = baselines.effectiveCoord(idx(ipart));
hlp.prep(coord.u, coord.v);
auto * RESTRICT wptr = hlp.p0w + u0*jump;
auto f0 = emb*emb;
if (have_wgt)
{
auto twgt = wgt(ipart);
f0*=twgt*twgt;
}
for (size_t cu=u0; cu<u1; ++cu)
{
auto f1=ku[cu]*ku[cu+du]*emb*emb;
auto f1=ku[cu]*ku[cu+du]*f0;
for (size_t cv=v0; cv<v1; ++cv)
wptr[cv] += f1*kv[cv]*kv[cv+dv];
wptr += jump;
......@@ -1592,33 +1523,25 @@ PYBIND11_MODULE(nifty_gridder, m)
"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, "grid_in"_a=None);
"idx"_a, "vis"_a, "grid_in"_a=None, "wgt"_a=None);
m.def("ms2grid",&ms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a,
"grid_in"_a=None);
m.def("ms2grid_wgt",&ms2grid_wgt<double>, "baselines"_a, "gconf"_a, "idx"_a,
"ms"_a, "wgt"_a, "grid_in"_a=None);
"grid_in"_a=None, "wgt"_a=None);
m.def("grid2vis",&grid2vis<double>, grid2vis_DS, "baselines"_a, "gconf"_a,
"idx"_a, "grid"_a);
"idx"_a, "grid"_a, "wgt"_a=None);
m.def("grid2ms",&grid2ms<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "ms_in"_a=None);
m.def("grid2ms_wgt",&grid2ms_wgt<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "wgt"_a, "ms_in"_a=None);
"grid"_a, "ms_in"_a=None, "wgt"_a=None);
m.def("vis2grid_c",&vis2grid_c<double>, vis2grid_c_DS, "baselines"_a,
"gconf"_a, "idx"_a, "vis"_a, "grid_in"_a=None);
"gconf"_a, "idx"_a, "vis"_a, "grid_in"_a=None, "wgt"_a=None);
m.def("ms2grid_c",&ms2grid_c<double>, ms2grid_c_DS, "baselines"_a, "gconf"_a,
"idx"_a, "ms"_a, "grid_in"_a=None);
m.def("ms2grid_c_wgt",&ms2grid_c_wgt<double>, "baselines"_a, "gconf"_a,
"idx"_a, "ms"_a, "wgt"_a, "grid_in"_a=None);
"idx"_a, "ms"_a, "grid_in"_a=None, "wgt"_a=None);
m.def("grid2vis_c",&grid2vis_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a);
"grid"_a, "wgt"_a=None);
m.def("grid2ms_c",&grid2ms_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "ms_in"_a=None);
m.def("grid2ms_c_wgt",&grid2ms_c_wgt<double>, "baselines"_a, "gconf"_a,
"idx"_a, "grid"_a, "wgt"_a, "ms_in"_a=None);
"grid"_a, "ms_in"_a=None, "wgt"_a=None);
m.def("apply_holo",&apply_holo<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a);
"grid"_a, "wgt"_a=None);
m.def("get_correlations", &get_correlations<double>, "baselines"_a, "gconf"_a,
"idx"_a, "du"_a, "dv"_a);
"idx"_a, "du"_a, "dv"_a, "wgt"_a=None);
m.def("set_nthreads",&set_nthreads, set_nthreads_DS, "nthreads"_a);
m.def("get_nthreads",&get_nthreads, get_nthreads_DS);
}
test.py
View file @ 50378509
......@@ -74,10 +74,9 @@ def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
assert id(grid2) == id(real_grid) # Same grid object
real_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.float64)
grid = ng.ms2grid_wgt(baselines, gconf, idx, ms,
weights, grid_in=None)
grid2 = ng.ms2grid_wgt(baselines, gconf, idx, ms,
weights, grid_in=real_grid)
grid = ng.ms2grid(baselines, gconf, idx, ms, grid_in=None, wgt=weights)
grid2 = ng.ms2grid(baselines, gconf, idx, ms, grid_in=real_grid,
wgt=weights)
assert_array_almost_equal(grid, grid2)
assert grid.dtype == real_grid.dtype == grid2.dtype
......@@ -108,10 +107,9 @@ def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
assert id(grid2) == id(cplx_grid) # Same grid object
cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
grid = ng.ms2grid_c_wgt(baselines, gconf, idx, ms,
weights, grid_in=None)
grid2 = ng.ms2grid_c_wgt(baselines, gconf, idx, ms,
weights, grid_in=cplx_grid)
grid = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=None, wgt=weights)
grid2 = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=cplx_grid,
wgt=weights)
# Almost same result
assert_array_almost_equal(grid, grid2)
......
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