Commit 9097edbb authored by Martin Reinecke's avatar Martin Reinecke
Browse files

add experimental get_diagonal() method

parent 1816b95e
......@@ -681,19 +681,19 @@ template<typename T> class GridderConfig
}
};
template<typename T> class Helper
template<typename T, typename T2=complex<T>> class Helper
{
private:
const GridderConfig<T> &gconf;
int nu, nv, nsafe, w;
T beta;
const complex<T> *grid_r;
complex<T> *grid_w;
const T2 *grid_r;
T2 *grid_w;
int su, sv;
int iu0, iv0; // start index of the current visibility
int bu0, bv0; // start index of the current buffer
vector<complex<T>> rbuf, wbuf;
vector<T2> rbuf, wbuf;
void dump() const
{
......@@ -733,12 +733,11 @@ template<typename T> class Helper
}
public:
const complex<T> *p0r;
complex<T> *p0w;
const T2 *p0r;
T2 *p0w;
vector<T> kernel;
Helper(const GridderConfig<T> &gconf_, const complex<T> *grid_r_,
complex<T> *grid_w_)
Helper(const GridderConfig<T> &gconf_, const T2 *grid_r_, T2 *grid_w_)
: gconf(gconf_), nu(gconf.Nu()), nv(gconf.Nv()), nsafe(gconf.Nsafe()),
w(gconf.W()), beta(gconf.Beta()), grid_r(grid_r_), grid_w(grid_w_),
su(2*nsafe+(1<<logsquare)), sv(2*nsafe+(1<<logsquare)),
......@@ -1274,6 +1273,53 @@ template<typename T> pyarr_c<complex<T>> apply_holo(
return res;
}
template<typename T> pyarr_c<T> get_diagonal(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, size_t dx, size_t dy)
{
size_t nu=gconf.Nu(), nv=gconf.Nv();
size_t w = gconf.W();
if (dx>=w) throw runtime_error("dx must be smaller than W");
if (dy>=w) throw runtime_error("dy must be smaller than W");
checkArray(idx_, "idx", {0});
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.template unchecked<1>();
auto res = makeArray<T>({nu, nv});
auto ogrid = res.mutable_data();
{
py::gil_scoped_release release;
T beta = gconf.Beta();
for (size_t i=0; i<nu*nv; ++i) ogrid[i] = 0.;
// Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
Helper<T,T> hlp(gconf, nullptr, ogrid);
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)
{
UVW<T> coord = baselines.effectiveCoord(idx(ipart));
hlp.prep(coord.u, coord.v);
auto * RESTRICT wptr = hlp.p0w;
for (size_t cu=0; cu<w-dx; ++cu)
{
auto f1=ku[cu]*ku[cu+dx]*emb*emb;
for (size_t cv=0; cv<w-dy; ++cv)
wptr[cv] += f1*kv[cv]*kv[cv+dy];
wptr += jump;
}
}
}
}
return res;
}
constexpr auto getIndices_DS = R"""(
Selects a subset of entries from a `Baselines` object.
......@@ -1438,6 +1484,8 @@ PYBIND11_MODULE(nifty_gridder, m)
"idx"_a, "grid"_a, "wgt"_a, "ms_in"_a=None);
m.def("apply_holo",&apply_holo<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a);
m.def("get_diagonal", &get_diagonal<double>, "baselines"_a, "gconf"_a,
"idx"_a, "dx"_a, "dy"_a);
m.def("set_nthreads",&set_nthreads, set_nthreads_DS, "nthreads"_a);
m.def("get_nthreads",&get_nthreads, get_nthreads_DS);
}
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