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

Merge branch 'wstacking' into 'master'

Wstacking

See merge request !7
parents 48289b09 d62f5f69
......@@ -412,6 +412,24 @@ template<typename T> class Baselines
np.array((nvis,), dtype=np.complex)
The visibility data for the index array
)""";
pyarr_c<T> ind2effectiveuvw(const pyarr_c<uint32_t> &idx_) const
{
checkArray(idx_, "idx", {0});
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.template unchecked<1>();
auto res_=makeArray<T>({nvis, 3});
auto res = res_.template mutable_unchecked<2>();
for (size_t i=0; i<nvis; i++)
{
auto uvw = effectiveCoord(idx(i));
res(i,0) = uvw.u;
res(i,1) = uvw.v;
res(i,2) = uvw.w;
}
return res_;
}
template<typename T2> pyarr_c<T2> ms2vis(const pyarr<T2> &ms_,
const pyarr_c<uint32_t> &idx_) const
{
......@@ -526,6 +544,24 @@ np.array((nxdirty, nydirty), dtype=np.float64)
the image with the taper applied
)""";
constexpr auto apply_wscreen_DS = R"""(
Applies the w screen to an image
Parameters
==========
dirty: nd.array((nxdirty, nydirty), dtype=np.complex128)
the image
w : float
the w value to use
adjoint: bool
if True, apply the complex conjugate of the w screen
Returns
=======
np.array((nxdirty, nydirty), dtype=np.complex128)
the image with the w screen applied
)""";
constexpr auto GridderConfig_DS = R"""(
Class storing information related to the gridding/degridding process.
......@@ -552,6 +588,15 @@ template<typename T> class GridderConfig
T beta;
vector<T> cfu, cfv;
complex<T> wscreen(double x, double y, double w, bool adjoint) const
{
constexpr double pi = 3.141592653589793238462643383279502884197;
double n = cos(sqrt(x+y)), xn = 1./n;
double phase = 2*pi*w*(n-1);
if (adjoint) phase *= -1;
return complex<T>(cos(phase)*xn, sin(phase)*xn);
}
public:
GridderConfig(size_t nxdirty, size_t nydirty, double epsilon,
double pixsize_x, double pixsize_y)
......@@ -716,6 +761,43 @@ template<typename T> class GridderConfig
iv0 = int(v-w*0.5 + 1 + nv) - nv;
if (iv0+w>nv+nsafe) iv0 = nv+nsafe-w;
}
pyarr_c<complex<T>> apply_wscreen(const pyarr_c<complex<T>> &dirty_, double w, bool adjoint) const
{
checkArray(dirty_, "dirty", {nx_dirty, ny_dirty});
auto dirty = dirty_.data();
auto res_ = makeArray<complex<T>>({nx_dirty, ny_dirty});
auto res = res_.mutable_data();
double x0 = -0.5*nx_dirty*psx,
y0 = -0.5*ny_dirty*psy;
{
py::gil_scoped_release release;
#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;
auto ws = wscreen(fx, fy*fy, w, adjoint);
res[ny_dirty*i+j] = dirty[ny_dirty*i+j]*ws; // lower left
size_t i2 = nx_dirty-i, j2 = ny_dirty-j;
if ((i>0)&&(i<i2))
{
res[ny_dirty*i2+j] = dirty[ny_dirty*i2+j]*ws; // lower right
if ((j>0)&&(j<j2))
res[ny_dirty*i2+j2] = dirty[ny_dirty*i2+j2]*ws; // upper right
}
if ((j>0)&&(j<j2))
res[ny_dirty*i+j2] = dirty[ny_dirty*i+j2]*ws; // upper left
}
}
}
}
return res_;
}
};
template<typename T, typename T2=complex<T>> class Helper
......@@ -1464,6 +1546,7 @@ PYBIND11_MODULE(nifty_gridder, m)
.def ("Nchannels",&Baselines<double>::Nchannels)
.def ("ms2vis",&Baselines<double>::ms2vis<complex<double>>,
Baselines<double>::ms2vis_DS, "ms"_a, "idx"_a)
.def ("ind2effectiveuvw",&Baselines<double>::ind2effectiveuvw, "idx"_a)
.def ("vis2ms",&Baselines<double>::vis2ms<complex<double>>,
Baselines<double>::vis2ms_DS, "vis"_a, "idx"_a, "ms_in"_a=None);
py::class_<GridderConfig<double>> (m, "GridderConfig", GridderConfig_DS)
......@@ -1484,6 +1567,8 @@ PYBIND11_MODULE(nifty_gridder, m)
.def("dirty2grid", &GridderConfig<double>::dirty2grid,
dirty2grid_DS, "dirty"_a)
.def("dirty2grid_c", &GridderConfig<double>::dirty2grid_c, "dirty"_a)
.def("apply_wscreen", &GridderConfig<double>::apply_wscreen,
apply_wscreen_DS, "dirty"_a, "w"_a, "adjoint"_a)
// pickle support
.def(py::pickle(
......
import nifty_gridder as ng
from time import time
import numpy as np
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)
wscreen = np.exp(2*np.pi*1j*w*(n - 1))/n
return wscreen
def time_op(func, x, ntries=5):
t0 = time()
for ii in range(ntries):
print(ii)
func(x)
return (time() - t0)/ntries
if __name__ == '__main__':
ng.set_nthreads(4)
ntries = 20
nx = 2048
dx = 12
w = 1000.2
ny, dy = nx, dx
conf = ng.GridderConfig(nx, ny, 1e-7, dx, dy)
x, y = conf.Nxdirty(), conf.Nydirty()
fld = np.random.randn(x, y) + 1j*np.random.randn(x, y)
func0 = lambda x: _wscreen(nx, dy, w).conjugate()*x
func1 = lambda x: conf.apply_wscreen(x, w, True)
print(time_op(func0, fld), time_op(func1, fld))
fld = np.random.randn(nx, ny)
func0 = lambda x: _wscreen(nx, dx, w)*x.real
func1 = lambda x: conf.apply_wscreen(x, w, False)
print(time_op(func0, fld), time_op(func1, fld))
......@@ -155,3 +155,28 @@ def test_pickling():
assert_(gc.Pixsize_y() == gc2.Pixsize_y())
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)
wscreen = np.exp(2*np.pi*1j*w*(n - 1))/n
return wscreen
@pmp('nx', [4, 18, 54])
@pmp('dx', [1., 0.13, 132])
@pmp('w', [0, 10, 8489])
def test_wstacking(nx, dx, w):
ny, dy = nx, dx
conf = ng.GridderConfig(nx, ny, 1e-7, dx, dy)
x, y = conf.Nu(), conf.Nv()
fld = np.random.randn(x, y) + 1j*np.random.randn(x, y)
res0 = (_wscreen(nx, dy, w).conjugate()*conf.grid2dirty_c(fld)).real
res1 = conf.apply_wscreen(conf.grid2dirty_c(fld), w, True).real
np.testing.assert_allclose(res0, res1)
fld = np.random.randn(nx, ny)
res0 = conf.dirty2grid_c(_wscreen(nx, dx, w)*fld.real)
res1 = conf.dirty2grid_c(conf.apply_wscreen(fld, w, False))
np.testing.assert_allclose(res0, res1)
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