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

first try

parent fc1b38de
...@@ -195,7 +195,7 @@ void checkArray(const py::array &arr, const char *aname, ...@@ -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) const vector<size_t> &shape)
{ {
if (in.is_none()) if (in.is_none())
...@@ -212,6 +212,16 @@ template<typename T> pyarr<T> provideArray(py::object &in, ...@@ -212,6 +212,16 @@ template<typename T> pyarr<T> provideArray(py::object &in,
return tmp_; 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, template<typename T> pyarr_c<T> provideCArray(py::object &in,
const vector<size_t> &shape) const vector<size_t> &shape)
{ {
...@@ -914,6 +924,8 @@ vis: np.array((nvis,), dtype=np.complex) ...@@ -914,6 +924,8 @@ vis: np.array((nvis,), dtype=np.complex)
The visibility data for the index array The visibility data for the index array
grid_in: np.array((nu,nv), dtype=np.complex128), optional grid_in: np.array((nu,nv), dtype=np.complex128), optional
If present, the result is added to this array. 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 Returns
======= =======
...@@ -923,13 +935,16 @@ np.array((nu,nv), dtype=np.complex128): ...@@ -923,13 +935,16 @@ np.array((nu,nv), dtype=np.complex128):
template<typename T> pyarr_c<complex<T>> vis2grid_c( template<typename T> pyarr_c<complex<T>> vis2grid_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf, const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &vis_, 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}); checkArray(vis_, "vis", {0});
size_t nvis = size_t(vis_.shape(0)); size_t nvis = size_t(vis_.shape(0));
checkArray(idx_, "idx", {nvis}); checkArray(idx_, "idx", {nvis});
auto vis=vis_.template unchecked<1>(); auto vis=vis_.template unchecked<1>();
auto idx = idx_.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(); size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = provideCArray<complex<T>>(grid_in, {nu, nv}); auto res = provideCArray<complex<T>>(grid_in, {nu, nv});
...@@ -955,6 +970,8 @@ template<typename T> pyarr_c<complex<T>> vis2grid_c( ...@@ -955,6 +970,8 @@ template<typename T> pyarr_c<complex<T>> vis2grid_c(
hlp.prep(coord.u, coord.v); hlp.prep(coord.u, coord.v);
auto * RESTRICT ptr = hlp.p0w; auto * RESTRICT ptr = hlp.p0w;
auto v(vis(ipart)*emb); auto v(vis(ipart)*emb);
if (have_wgt)
v*=wgt(ipart);
for (size_t cu=0; cu<w; ++cu) for (size_t cu=0; cu<w; ++cu)
{ {
complex<T> tmp(v*ku[cu]); complex<T> tmp(v*ku[cu]);
...@@ -984,6 +1001,8 @@ vis: np.array((nvis,), dtype=np.complex) ...@@ -984,6 +1001,8 @@ vis: np.array((nvis,), dtype=np.complex)
The visibility data for the index array The visibility data for the index array
grid_in: np.array((nu,nv), dtype=np.float64), optional grid_in: np.array((nu,nv), dtype=np.float64), optional
If present, the result is added to this array. 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 Returns
======= =======
...@@ -992,8 +1011,8 @@ np.array((nu,nv), dtype=np.float64): ...@@ -992,8 +1011,8 @@ np.array((nu,nv), dtype=np.float64):
)"""; )""";
template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines, template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_, const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr<complex<T>> &vis_, py::object &grid_in) const pyarr<complex<T>> &vis_, py::object &grid_in, const py::object &wgt_)
{ return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_, None), grid_in); } { return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_, None, wgt_), grid_in); }
constexpr auto ms2grid_c_DS = R"""( constexpr auto ms2grid_c_DS = R"""(
Grids measurement set data onto a UV grid Grids measurement set data onto a UV grid
...@@ -1011,6 +1030,8 @@ ms: np.array((nrows, nchannels), dtype=np.complex128) ...@@ -1011,6 +1030,8 @@ ms: np.array((nrows, nchannels), dtype=np.complex128)
the measurement set. the measurement set.
grid_in: np.array((nu,nv), dtype=np.complex128), optional grid_in: np.array((nu,nv), dtype=np.complex128), optional
If present, the result is added to this array. 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 Returns
======= =======
...@@ -1020,7 +1041,7 @@ np.array((nu,nv), dtype=np.complex128): ...@@ -1020,7 +1041,7 @@ np.array((nu,nv), dtype=np.complex128):
template<typename T> pyarr_c<complex<T>> ms2grid_c( template<typename T> pyarr_c<complex<T>> ms2grid_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf, const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_, 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 nrows = baselines.Nrows();
auto nchan = baselines.Nchannels(); auto nchan = baselines.Nchannels();
...@@ -1029,6 +1050,9 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c( ...@@ -1029,6 +1050,9 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c(
size_t nvis = size_t(idx_.shape(0)); size_t nvis = size_t(idx_.shape(0));
auto ms = ms_.template unchecked<2>(); auto ms = ms_.template unchecked<2>();
auto idx = idx_.template unchecked<1>(); 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(); size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = provideCArray<complex<T>>(grid_in, {nu, nv}); auto res = provideCArray<complex<T>>(grid_in, {nu, nv});
...@@ -1057,6 +1081,8 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c( ...@@ -1057,6 +1081,8 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c(
hlp.prep(coord.u, coord.v); hlp.prep(coord.u, coord.v);
auto * RESTRICT ptr = hlp.p0w; auto * RESTRICT ptr = hlp.p0w;
auto v(ms(row,chan)*emb); auto v(ms(row,chan)*emb);
if (have_wgt)
v*=wgt(row, chan);
for (size_t cu=0; cu<w; ++cu) for (size_t cu=0; cu<w; ++cu)
{ {
complex<T> tmp(v*ku[cu]); complex<T> tmp(v*ku[cu]);
...@@ -1072,73 +1098,13 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c( ...@@ -1072,73 +1098,13 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c(
template<typename T> pyarr_c<T> ms2grid(const Baselines<T> &baselines, template<typename T> pyarr_c<T> ms2grid(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_, const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr<complex<T>> &ms_, py::object &grid_in) const pyarr<complex<T>> &ms_, py::object &grid_in, const py::object &wgt_)
{ return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_, None), grid_in); } { return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_, None, wgt_), 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); }
template<typename T> pyarr_c<complex<T>> grid2vis_c( template<typename T> pyarr_c<complex<T>> grid2vis_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf, 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(); size_t nu=gconf.Nu(), nv=gconf.Nv();
checkArray(idx_, "idx", {0}); checkArray(idx_, "idx", {0});
...@@ -1146,6 +1112,9 @@ template<typename T> pyarr_c<complex<T>> grid2vis_c( ...@@ -1146,6 +1112,9 @@ template<typename T> pyarr_c<complex<T>> grid2vis_c(
checkArray(grid_, "grid", {nu, nv}); checkArray(grid_, "grid", {nu, nv});
size_t nvis = size_t(idx_.shape(0)); size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.template unchecked<1>(); 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 res = makeArray<complex<T>>({nvis});
auto vis = res.mutable_data(); auto vis = res.mutable_data();
...@@ -1178,6 +1147,7 @@ template<typename T> pyarr_c<complex<T>> grid2vis_c( ...@@ -1178,6 +1147,7 @@ template<typename T> pyarr_c<complex<T>> grid2vis_c(
r += tmp*ku[cu]; r += tmp*ku[cu];
ptr += jump; ptr += jump;
} }
if (have_wgt) r*=wgt[ipart];
vis[ipart] = r*emb; vis[ipart] = r*emb;
} }
} }
...@@ -1199,6 +1169,8 @@ idx: np.array((nvis,), dtype=np.uint32) ...@@ -1199,6 +1169,8 @@ idx: np.array((nvis,), dtype=np.uint32)
the indices for the entries to be degridded the indices for the entries to be degridded
grid: np.array((nu,nv), dtype=np.float64): grid: np.array((nu,nv), dtype=np.float64):
the gridded visibilities (made real by making use of Hermitian symmetry) 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 Returns
======= =======
...@@ -1207,12 +1179,12 @@ np.array((nvis,), dtype=np.complex) ...@@ -1207,12 +1179,12 @@ np.array((nvis,), dtype=np.complex)
)"""; )""";
template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines, template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_, const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr_c<T> &grid_) const pyarr_c<T> &grid_, const py::object &wgt_)
{ return grid2vis_c(baselines, gconf, idx_, hartley2complex(grid_)); } { return grid2vis_c(baselines, gconf, idx_, hartley2complex(grid_), wgt_); }
template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines, template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_, 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 nrows = baselines.Nrows();
auto nchan = baselines.Nchannels(); auto nchan = baselines.Nchannels();
...@@ -1222,6 +1194,9 @@ template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines ...@@ -1222,6 +1194,9 @@ template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines
checkArray(grid_, "grid", {nu, nv}); checkArray(grid_, "grid", {nu, nv});
size_t nvis = size_t(idx_.shape(0)); size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.template unchecked<1>(); 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 res = provideArray<complex<T>>(ms_in, {nrows, nchan});
auto ms = res.template mutable_unchecked<2>(); auto ms = res.template mutable_unchecked<2>();
...@@ -1257,6 +1232,9 @@ template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines ...@@ -1257,6 +1232,9 @@ template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines
r += tmp*ku[cu]; r += tmp*ku[cu];
ptr += jump; ptr += jump;
} }
r*=emb;
if (have_wgt)
r*=wgt(row, chan);
ms(row,chan) += r*emb; ms(row,chan) += r*emb;
} }
} }
...@@ -1266,78 +1244,13 @@ template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines ...@@ -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, template<typename T> pyarr_c<complex<T>> grid2ms(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_, const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr_c<T> &grid_, py::object &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); } { return grid2ms_c(baselines, gconf, idx_, hartley2complex(grid_), ms_in, wgt_); }
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);
}
template<typename T> pyarr_c<complex<T>> apply_holo( template<typename T> pyarr_c<complex<T>> apply_holo(
const Baselines<T> &baselines, const GridderConfig<T> &gconf, 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(); size_t nu=gconf.Nu(), nv=gconf.Nv();
checkArray(idx_, "idx", {0}); checkArray(idx_, "idx", {0});
...@@ -1345,6 +1258,9 @@ template<typename T> pyarr_c<complex<T>> apply_holo( ...@@ -1345,6 +1258,9 @@ template<typename T> pyarr_c<complex<T>> apply_holo(
checkArray(grid_, "grid", {nu, nv}); checkArray(grid_, "grid", {nu, nv});
size_t nvis = size_t(idx_.shape(0)); size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.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>();
auto res = makeArray<complex<T>>({nu, nv}); auto res = makeArray<complex<T>>({nu, nv});
auto ogrid = res.mutable_data(); auto ogrid = res.mutable_data();
...@@ -1365,7 +1281,8 @@ template<typename T> pyarr_c<complex<T>> apply_holo( ...@@ -1365,7 +1281,8 @@ template<typename T> pyarr_c<complex<T>> apply_holo(
#pragma omp for schedule(guided,100) #pragma omp for schedule(guided,100)
for (size_t ipart=0; ipart<nvis; ++ipart) 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); hlp.prep(coord.u, coord.v);
complex<T> r = 0; complex<T> r = 0;
const auto * RESTRICT ptr = hlp.p0r; const auto * RESTRICT ptr = hlp.p0r;
...@@ -1378,6 +1295,11 @@ template<typename T> pyarr_c<complex<T>> apply_holo( ...@@ -1378,6 +1295,11 @@ template<typename T> pyarr_c<complex<T>> apply_holo(
ptr += jump; ptr += jump;
} }
r*=emb*emb; r*=emb*emb;
if (have_wgt)
{
auto twgt = wgt(ipart);
r*=twgt*twgt;
}
auto * RESTRICT wptr = hlp.p0w; auto * RESTRICT wptr = hlp.p0w;
for (size_t cu=0; cu<w; ++cu) for (size_t cu=0; cu<w; ++cu)
{ {
...@@ -1394,7 +1316,7 @@ template<typename T> pyarr_c<complex<T>> apply_holo( ...@@ -1394,7 +1316,7 @@ template<typename T> pyarr_c<complex<T>> apply_holo(
template<typename T> pyarr_c<T> get_correlations( template<typename T> pyarr_c<T> get_correlations(
const Baselines<T> &baselines, const GridderConfig<T> &gconf, 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 nu=gconf.Nu(), nv=gconf.Nv();
size_t w = gconf.W(); size_t w = gconf.W();
...@@ -1403,6 +1325,9 @@ template<typename T> pyarr_c<T> get_correlations( ...@@ -1403,6 +1325,9 @@ template<typename T> pyarr_c<T> get_correlations(
checkArray(idx_, "idx", {0}); checkArray(idx_, "idx", {0});
size_t nvis = size_t(idx_.shape(0)); size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.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>();
auto res = makeArray<T>({nu, nv}); auto res = makeArray<T>({nu, nv});
auto ogrid = res.mutable_data(); auto ogrid = res.mutable_data();
...@@ -1436,9 +1361,15 @@ template<typename T> pyarr_c<T> get_correlations( ...@@ -1436,9 +1361,15 @@ template<typename T> pyarr_c<T> get_correlations(
UVW<T> coord = baselines.effectiveCoord(idx(ipart)); UVW<T> coord = baselines.effectiveCoord(idx(ipart));
hlp.prep(coord.u, coord.v); hlp.prep(coord.u, coord.v);
auto * RESTRICT wptr = hlp.p0w + u0*jump; 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) 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) for (size_t cv=v0; cv<v1; ++cv)
wptr[cv] += f1*kv[cv]*kv[cv+dv]; wptr[cv] += f1*kv[cv]*kv[cv+dv];
wptr += jump; wptr += jump;
...@@ -1592,33 +1523,25 @@ PYBIND11_MODULE(nifty_gridder, m) ...@@ -1592,33 +1523,25 @@ PYBIND11_MODULE(nifty_gridder, m)
"gconf"_a, "flags"_a, "chbegin"_a=-1, "chend"_a=-1, "gconf"_a, "flags"_a, "chbegin"_a=-1, "chend"_a=-1,
"wmin"_a=-1e30, "wmax"_a=1e30); "wmin"_a=-1e30, "wmax"_a=1e30);
m.def("vis2grid",&vis2grid<double>, vis2grid_DS, "baselines"_a, "gconf"_a, 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, m.def("ms2grid",&ms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a,
"grid_in"_a=None); "grid_in"_a=None, "wgt"_a=None);
m.def("ms2grid_wgt",&ms2grid_wgt<double>, "baselines"_a, "gconf"_a, "idx"_a,
"ms"_a, "wgt"_a, "grid_in"_a=None);
m.def("grid2vis",&grid2vis<double>, grid2vis_DS, "baselines"_a, "gconf"_a, 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, m.def("grid2ms",&grid2ms<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "ms_in"_a=None); "grid"_a, "ms_in"_a=None, "wgt"_a=None);
m.def("grid2ms_wgt",&grid2ms_wgt<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "wgt"_a, "ms_in"_a=None);
m.def("vis2grid_c",&vis2grid_c<double>, vis2grid_c_DS, "baselines"_a, 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, m.def("ms2grid_c",&ms2grid_c<double>, ms2grid_c_DS, "baselines"_a, "gconf"_a,
"idx"_a, "ms"_a, "grid_in"_a=None); "idx"_a, "ms"_a, "grid_in"_a=None, "wgt"_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);
m.def("grid2vis_c",&grid2vis_c<double>, "baselines"_a, "gconf"_a, "idx"_a, 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, m.def("grid2ms_c",&grid2ms_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "ms_in"_a=None); "grid"_a, "ms_in"_a=None, "wgt"_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);
m.def("apply_holo",&apply_holo<double>, "baselines"_a, "gconf"_a, "idx"_a, 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, 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("set_nthreads",&set_nthreads, set_nthreads_DS, "nthreads"_a);
m.def("get_nthreads",&get_nthreads, get_nthreads_DS); m.def("get_nthreads",&get_nthreads, get_nthreads_DS);
} }