Skip to content

GitLab

Commit 73b99e5c authored by Martin Reinecke's avatar Martin Reinecke
Browse files

more functions

parent 6a2c68ec
Hide whitespace changes
Inline Side-by-side
gridder_cxx.h
View file @ 73b99e5c
......@@ -108,6 +108,26 @@ template<typename T, size_t ndim> class mav
template<typename T, size_t ndim> using const_mav = mav<const T, ndim>;
template<typename T, size_t ndim> class tmpStorage
{
private:
vector<T> d;
mav<T,ndim> mav_;
static size_t prod(const array<size_t,ndim> &shp)
{
size_t res=1;
for (auto v: shp) res*=v;
return res;
}
public:
tmpStorage(const array<size_t,ndim> &shp)
: d(prod(shp)), mav_(d.data(), shp) {}
mav<T,ndim> &getMav() { return mav_; };
void fill(const T & val)
{ std::fill(d.begin(), d.end(), val); }
};
//
// basic utilities
//
......@@ -384,13 +404,7 @@ template<typename T> class GridderConfig
complex<T> wscreen(double x, double y, double w, bool adjoint) const
{
constexpr double pi = 3.141592653589793238462643383279502884197;
#if 1
double eps = sqrt(x+y);
double s = sin(eps);
double nm1 = -s*s/(1.+cos(eps));
#else
double nm1 = (-x-y)/(sqrt(1.-x-y)+1);
#endif
double nm1 = (-x-y)/(sqrt(1.-x-y)+1); // more accurate form of sqrt(1-x-y)
double n = nm1+1., xn = 1./n;
double phase = 2*pi*w*nm1;
if (adjoint) phase *= -1;
......@@ -827,6 +841,164 @@ template<typename T> void grid2ms_c
}
}
template<typename T> void apply_holo
(const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,2> &grid,
mav<complex<T>,2> &ogrid, const const_mav<T,1> &wgt)
{
size_t nu=gconf.Nu(), nv=gconf.Nv();
checkShape(grid.shape(), {nu, nv});
size_t nvis = idx.shape(0);
bool have_wgt = wgt.size()!=0;
if (have_wgt) checkShape(wgt.shape(), {nvis});
checkShape(ogrid.shape(), {nu, nv});
for (size_t i=0; i<nu; ++i)
for (size_t j=0; j<nv; ++j)
ogrid(i,j) = T(0);
T beta = gconf.Beta();
size_t supp = gconf.Supp();
// Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
Helper<T> hlp(gconf, grid.data(), ogrid.data());
T emb = exp(-2*beta);
int jump = hlp.lineJump();
const T * ku = hlp.kernel.data();
const T * kv = hlp.kernel.data()+supp;
#pragma omp for schedule(guided,100)
for (size_t ipart=0; ipart<nvis; ++ipart)
{
auto tidx = idx(ipart);
UVW<T> coord = baselines.effectiveCoord(tidx);
hlp.prep(coord.u, coord.v);
complex<T> r = 0;
const auto * ptr = hlp.p0r;
for (size_t cu=0; cu<supp; ++cu)
{
complex<T> tmp(0);
for (size_t cv=0; cv<supp; ++cv)
tmp += ptr[cv] * kv[cv];
r += tmp*ku[cu];
ptr += jump;
}
r*=emb*emb;
if (have_wgt)
{
auto twgt = wgt(ipart);
r*=twgt*twgt;
}
auto * wptr = hlp.p0w;
for (size_t cu=0; cu<supp; ++cu)
{
complex<T> tmp(r*ku[cu]);
for (size_t cv=0; cv<supp; ++cv)
wptr[cv] += tmp*kv[cv];
wptr += jump;
}
}
}
}
template<typename T> void get_correlations
(const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const const_mav<uint32_t,1> &idx, int du, int dv,
mav<T,2> &ogrid, const const_mav<T,1> &wgt)
{
size_t nu=gconf.Nu(), nv=gconf.Nv();
size_t nvis = idx.shape(0);
bool have_wgt = wgt.size()!=0;
if (have_wgt) checkShape(wgt.shape(), {nvis});
checkShape(ogrid.shape(), {nu, nv});
size_t supp = gconf.Supp();
myassert(size_t(abs(du))<supp, "|du| must be smaller than Supp");
myassert(size_t(abs(dv))<supp, "|dv| must be smaller than Supp");
T beta = gconf.Beta();
for (size_t i=0; i<nu; ++i)
for (size_t j=0; j<nv; ++j)
ogrid(i,j) = T(0);
size_t u0, u1, v0, v1;
if (du>=0)
{ u0=0; u1=supp-du; }
else
{ u0=-du; u1=supp; }
if (dv>=0)
{ v0=0; v1=supp-dv; }
else
{ v0=-dv; v1=supp; }
// Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
Helper<T,T> hlp(gconf, nullptr, ogrid.data());
T emb = exp(-2*beta);
int jump = hlp.lineJump();
const T * ku = hlp.kernel.data();
const T * kv = hlp.kernel.data()+supp;
#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 * 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]*f0;
for (size_t cv=v0; cv<v1; ++cv)
wptr[cv] += f1*kv[cv]*kv[cv+dv];
wptr += jump;
}
}
}
}
template<typename T> void apply_wcorr(const GridderConfig<T> &gconf,
mav<complex<T>,2> &dirty, const EC_Kernel_with_correction<T> &kernel, double dw)
{
auto nx_dirty=gconf.Nxdirty();
auto ny_dirty=gconf.Nydirty();
auto psx=gconf.Pixsize_x();
auto psy=gconf.Pixsize_y();
double x0 = -0.5*nx_dirty*psx,
y0 = -0.5*ny_dirty*psy;
#pragma omp parallel for schedule(static) num_threads(nthreads)
for (size_t i=0; i<=nx_dirty/2; ++i)
{
double fx = x0+i*psx;
fx *= fx;
for (size_t j=0; j<=ny_dirty/2; ++j)
{
double fy = y0+j*psy;
fy*=fy;
auto nm1 = (-fx-fy)/(sqrt(1.-fx-fy)+1.); // accurate form of sqrt(1-x-y)
T fct = kernel.corfac(nm1*dw);
size_t i2 = nx_dirty-i, j2 = ny_dirty-j;
dirty(i,j)*=fct;
if ((i>0)&&(i<i2))
{
dirty(i2,j)*=fct;
if ((j>0)&&(j<j2))
dirty(i2,j2)*=fct;
}
if ((j>0)&&(j<j2))
dirty(i,j2)*=fct;
}
}
}
template<typename T> void vis2dirty_wstack(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const const_mav<uint32_t,1> &idx,
const const_mav<complex<T>,1> &vis, mav<complex<T>,2> &dirty)
......@@ -857,8 +1029,7 @@ cout << "data w range: " << wmin << " to " << wmax << endl;
double nmin = sqrt(max(1.-x0*x0-y0*y0,0.))-1.;
double dw = 0.25/abs(nmin);
double w_eps=1e-7; // FIXME
auto w_supp = get_supp(w_eps);
auto w_supp = gconf.Supp();
EC_Kernel_with_correction<T> kernel(w_supp);
wmin -= 0.5*w_supp*dw;
wmax += 0.5*w_supp*dw;
......@@ -877,6 +1048,13 @@ cout << "nplanes: " << nplanes << endl;
for (size_t i=0; i<nx_dirty; ++i)
for (size_t j=0; j<ny_dirty; ++j)
dirty(i,j) = 0.;
{
vector<complex<T>> vis_loc_;
vector<uint32_t> idx_loc_;
tmpStorage<complex<T>,2> grid_({nu,nv});
auto grid=grid_.getMav();
tmpStorage<complex<T>,2> tdirty_({nx_dirty,ny_dirty});
auto tdirty=tdirty_.getMav();
for (size_t iw=0; iw<nplanes; ++iw)
{
cout << "working on w plane #" << iw << " containing " << nvis_plane[iw]
......@@ -884,64 +1062,37 @@ cout << "working on w plane #" << iw << " containing " << nvis_plane[iw]
if (nvis_plane[iw]==0) continue;
double wcur = wmin+iw*dw;
size_t cnt=0;
vector<complex<T>> vis_loc_ (nvis_plane[iw]);
cout << "blip0" << endl;
vis_loc_.resize(nvis_plane[iw]);
auto vis_loc = mav<complex<T>,1>(vis_loc_.data(), {nvis_plane[iw]});
vector<uint32_t> idx_loc_ (nvis_plane[iw]);
idx_loc_.resize(nvis_plane[iw]);
auto idx_loc = mav<uint32_t,1>(idx_loc_.data(), {nvis_plane[iw]});
for (size_t ipart=0; ipart<nvis; ++ipart)
if ((iw>=minplane[ipart]) && (iw<minplane[ipart]+w_supp))
{
myassert(cnt<nvis_plane[iw],"must not happen");
double x=min(1.,2./(w_supp*dw)*abs(wcur-wval[ipart]));
vis_loc[cnt] = vis[ipart]*kernel(x);
idx_loc[cnt] = idx[ipart];
++cnt;
}
myassert(cnt==nvis_plane[iw],"must not happen 2");
auto grid_=vector<complex<T>>(nu*nv);
auto grid=mav<complex<T>,2>(grid_.data(),{nu, nv});
cout << "blip1" << endl;
grid_.fill(0.);
vis2grid_c(baselines, gconf, const_mav<uint32_t,1>(idx_loc), const_mav<complex<T>,1>(vis_loc), grid, const_mav<T,1>(nullptr,{0}));
auto tdirty_=vector<complex<T>>(nx_dirty*ny_dirty);
auto tdirty=mav<complex<T>,2>(tdirty_.data(),{nx_dirty, ny_dirty});
cout << "blip2" << endl;
gconf.grid2dirty_c(grid, tdirty);
cout << "blip3" << endl;
gconf.apply_wscreen(tdirty, tdirty, wcur, false);
cout << "blip4" << endl;
for (size_t i=0; i<nx_dirty; ++i)
for (size_t j=0; j<ny_dirty; ++j)
dirty(i,j) += tdirty(i,j);
}
}
// correct for w gridding
cout << "applying correction for gridding in w direction" << endl;
#pragma omp parallel num_threads(nthreads)
{
#pragma omp for schedule(static)
for (size_t i=0; i<=nx_dirty/2; ++i)
{
double fx = x0+i*psx;
fx *= fx;
for (size_t j=0; j<=ny_dirty/2; ++j)
{
double fy = y0+j*psy;
fy*=fy;
#if 1
auto nm1=cos(sqrt(fx+fy)) -1; // cosine profile
#else
auto nm1 = (-fx-fy)/(sqrt(1.-fx-fy)+1.);
#endif
T fct = kernel.corfac(nm1*dw);
size_t i2 = nx_dirty-i, j2 = ny_dirty-j;
dirty(i,j)*=fct;
if ((i>0)&&(i<i2))
{
dirty(i2,j)*=fct;
if ((j>0)&&(j<j2))
dirty(i2,j2)*=fct;
}
if ((j>0)&&(j<j2))
dirty(i,j2)*=fct;
}
}
}
apply_wcorr(gconf, dirty, kernel, dw);
}
} // namespace gridder
......
nifty_gridder.cc
View file @ 73b99e5c
......@@ -561,137 +561,38 @@ template<typename T> pyarr<complex<T>> Pygrid2ms_c(
return res;
}
template<typename T> pyarr<complex<T>> apply_holo(
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_,
const py::object &wgt_)
{
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>();
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();
auto idx = make_const_mav<1>(idx_);
auto grid = make_const_mav<2>(grid_);
pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {idx.shape(0)});
auto wgt=make_const_mav<1>(wgt2);
auto res = makeArray<complex<T>>({grid.shape(0),grid.shape(1)});
auto ogrid = make_mav<2>(res);
{
py::gil_scoped_release release;
for (size_t i=0; i<nu*nv; ++i) ogrid[i] = T(0);
T beta = gconf.Beta();
size_t supp = gconf.Supp();
// Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
Helper<T> hlp(gconf, grid, ogrid);
T emb = exp(-2*beta);
int jump = hlp.lineJump();
const T * ku = hlp.kernel.data();
const T * kv = hlp.kernel.data()+supp;
#pragma omp for schedule(guided,100)
for (size_t ipart=0; ipart<nvis; ++ipart)
{
auto tidx = idx(ipart);
UVW<T> coord = baselines.effectiveCoord(tidx);
hlp.prep(coord.u, coord.v);
complex<T> r = 0;
const auto * ptr = hlp.p0r;
for (size_t cu=0; cu<supp; ++cu)
{
complex<T> tmp(0);
for (size_t cv=0; cv<supp; ++cv)
tmp += ptr[cv] * kv[cv];
r += tmp*ku[cu];
ptr += jump;
}
r*=emb*emb;
if (have_wgt)
{
auto twgt = wgt(ipart);
r*=twgt*twgt;
}
auto * wptr = hlp.p0w;
for (size_t cu=0; cu<supp; ++cu)
{
complex<T> tmp(r*ku[cu]);
for (size_t cv=0; cv<supp; ++cv)
wptr[cv] += tmp*kv[cv];
wptr += jump;
}
}
}
apply_holo(baselines, gconf, idx, grid, ogrid, wgt);
}
return res;
}
template<typename T> pyarr<T> get_correlations(
template<typename T> pyarr<T> Pyget_correlations(
const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, int du, int dv, const py::object &wgt_)
{
size_t nu=gconf.Nu(), nv=gconf.Nv();
size_t supp = gconf.Supp();
myassert(size_t(abs(du))<supp, "|du| must be smaller than Supp");
myassert(size_t(abs(dv))<supp, "|dv| must be smaller than Supp");
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();
auto idx = make_const_mav<1>(idx_);
pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {idx.shape(0)});
auto wgt=make_const_mav<1>(wgt2);
auto res = makeArray<T>({gconf.Nu(),gconf.Nv()});
auto ogrid = make_mav<2>(res);
{
py::gil_scoped_release release;
T beta = gconf.Beta();
for (size_t i=0; i<nu*nv; ++i) ogrid[i] = 0.;
size_t u0, u1, v0, v1;
if (du>=0)
{ u0=0; u1=supp-du; }
else
{ u0=-du; u1=supp; }
if (dv>=0)
{ v0=0; v1=supp-dv; }
else
{ v0=-dv; v1=supp; }
// 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 * ku = hlp.kernel.data();
const T * kv = hlp.kernel.data()+supp;
#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 * 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]*f0;
for (size_t cv=v0; cv<v1; ++cv)
wptr[cv] += f1*kv[cv]*kv[cv+dv];
wptr += jump;
}
}
}
get_correlations(baselines, gconf, idx, du, dv, ogrid, wgt);
}
return res;
}
......@@ -860,9 +761,9 @@ PYBIND11_MODULE(nifty_gridder, m)
"grid"_a, "wgt"_a=None);
m.def("grid2ms_c",&Pygrid2ms_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "ms_in"_a=None, "wgt"_a=None);
m.def("apply_holo",&apply_holo<double>, "baselines"_a, "gconf"_a, "idx"_a,
m.def("apply_holo",&Pyapply_holo<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "wgt"_a=None);
m.def("get_correlations", &get_correlations<double>, "baselines"_a, "gconf"_a,
m.def("get_correlations", &Pyget_correlations<double>, "baselines"_a, "gconf"_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 @ 73b99e5c
......@@ -25,8 +25,7 @@ def _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow):
def _wscreen(npix, dst, w):
dc = (np.linspace(start=-npix/2, stop=npix/2 - 1, num=npix)*dst)**2
ls = np.broadcast_to(dc, (dc.shape[0],)*2)
theta = np.sqrt(ls + ls.T)
n = np.cos(theta)
n = np.sqrt(1-ls-ls.T)
wscreen = np.exp(2*np.pi*1j*w*(n - 1))/n
return wscreen
......@@ -274,15 +273,12 @@ def test_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
y *= conf.Pixsize_y()
res1 = np.zeros_like(res0)
eps = np.sqrt(x**2+y**2)
s = np.sin(eps)
nm1 = -s*s/(1+np.cos(eps))
n = nm1+1
n=np.sqrt(1-x**2-y**2)
nm1=n-1
for ii in idx:
phase = x*uvw[ii, 0] + y*uvw[ii, 1] + uvw[ii, 2]*nm1
res1 += (vis[ii]*np.exp(2j*np.pi*phase)).real
res1 /= n
print(_l2error(res0, res1))
assert_(_l2error(res0, res1) < 1e-4)
@pmp('nxdirty', [16, 64])
......@@ -325,13 +321,10 @@ def test_against_wdft_exact(nxdirty, nydirty, nchan, nrow, fov):
y *= conf.Pixsize_y()
res1 = np.zeros_like(res0)
eps = np.sqrt(x**2+y**2)
s = np.sin(eps)
nm1 = -s*s/(1+np.cos(eps))
n = nm1+1
n = np.sqrt(1-x**2-y**2)
nm1 = n-1
for ii in idx:
phase = x*uvw[ii, 0] + y*uvw[ii, 1] + uvw[ii, 2]*nm1
res1 += (vis[ii]*np.exp(2j*np.pi*phase)).real
res1 /= n
print(_l2error(res0, res1))
assert_(_l2error(res0, res1) < 1e-4)
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