Fix C/Python interface

Remove the grid argument from transform and add set_grid and get_grid

idg-api/BufferSet.cpp

@@ -268,8 +268,8 @@ void BufferSetImpl::init(size_t size, float cell_size, float max_w,
 
   m_grid.reset(new Grid(nr_w_layers, 4, m_padded_size, m_padded_size));
   m_grid->zero();
-  m_proxy->set_grid(m_grid, m_subgridsize, m_image_size, m_w_step,
-                    m_shift.data());
+  m_proxy->set_grid(m_grid);
+  m_proxy->init_wtiles(m_subgridsize);
 
   m_taper_subgrid.resize(m_subgridsize);
   m_taper_grid.resize(m_padded_size);
@@ -585,6 +585,7 @@ void BufferSetImpl::get_image(double* image) {
   m_get_image_watch->Start();
 
   // Flush all pending operations on the grid
+  m_proxy->flush_wtiles();
   m_proxy->get_grid();
 
   double runtime = -omp_get_wtime();

idg-bin/examples/cxx/common/common.h

@@ -249,8 +249,7 @@ void run() {
   // Iterate all cycles
   for (unsigned cycle = 0; cycle < nr_cycles; cycle++) {
     // Set grid
-    float w_step = 0.0;
-    proxy.set_grid(grid, subgrid_size, image_size, w_step, shift.data());
+    proxy.set_grid(grid);
 
     // Iterate all time blocks
     for (unsigned time_offset = 0; time_offset < total_nr_timesteps;

idg-bin/examples/python/common.py

@@ -84,9 +84,11 @@ def gridding(
     p.gridding(
         w_step, shift, cell_size, kernel_size, subgrid_size,
         frequencies, visibilities, uvw, baselines,
-        grid, aterms, aterms_offsets, spheroidal)
+        aterms, aterms_offsets, spheroidal)
+    p.get_grid(grid)
     util.plot_grid(grid, scaling='log')
-    p.transform(idg.FourierDomainToImageDomain, grid)
+    p.transform(idg.FourierDomainToImageDomain)
+    p.get_grid(grid)
     util.plot_grid(grid)
     #util.plot_grid(grid, pol=0)
 
@@ -97,11 +99,11 @@ def gridding(
 def degridding(
         p, w_step, shift, cell_size, kernel_size, subgrid_size, frequencies, visibilities,
         uvw, baselines, grid, aterms, aterms_offsets, spheroidal):
-    p.transform(idg.ImageDomainToFourierDomain, grid)
+    p.transform(idg.ImageDomainToFourierDomain)
     p.degridding(
         w_step, shift, cell_size, kernel_size, subgrid_size,
         frequencies, visibilities, uvw, baselines,
-        grid, aterms, aterms_offsets, spheroidal)
+        aterms, aterms_offsets, spheroidal)
     #util.plot_visibilities(visibilities)
 
 

idg-bin/tests/cxx/CPU/Reference/main.cpp

@@ -103,8 +103,7 @@ int test01() {
   idg::proxy::cpu::Reference proxy;
 
   // Set grid
-  float w_step = 0.0;
-  proxy.set_grid(grid, subgrid_size, image_size, w_step, shift.data());
+  proxy.set_grid(grid);
 
   // Create plan
   clog << ">>> Create plan" << endl;
@@ -132,7 +131,7 @@ int test01() {
   proxy.transform(idg::ImageDomainToFourierDomain);
 
   // Set reference grid
-  proxy.set_grid(grid_ref, subgrid_size, image_size, w_step, shift.data());
+  proxy.set_grid(grid_ref);
 
   proxy.degridding(plan, w_offset, shift, cell_size, kernel_size, subgrid_size,
                    frequencies, visibilities, uvw, baselines, aterms,

idg-bin/tests/cxx/common/common.h

@@ -178,9 +178,8 @@ int compare_to_reference(float tol = 1000 *
   }
 
   // Bind the grids to the respective proxies
-  float w_step = 0.0;
-  optimized.set_grid(grid, subgrid_size, image_size, w_step, shift.data());
-  reference.set_grid(grid_ref, subgrid_size, image_size, w_step, shift.data());
+  optimized.set_grid(grid);
+  reference.set_grid(grid_ref);
 
   // Set w-terms to zero
   for (unsigned bl = 0; bl < nr_baselines; bl++) {
@@ -200,14 +199,14 @@ int compare_to_reference(float tol = 1000 *
 #if TEST_GRIDDING
   // Run gridder
   std::clog << ">>> Run gridding" << std::endl;
-  optimized.set_grid(grid, subgrid_size, image_size, w_step, shift.data());
+  optimized.set_grid(grid);
   optimized.gridding(plan, w_offset, shift, cell_size, kernel_size,
                      subgrid_size, frequencies, visibilities, uvw, baselines,
                      aterms, aterms_offsets, spheroidal);
   optimized.get_grid();
 
   std::clog << ">>> Run reference gridding" << std::endl;
-  reference.set_grid(grid_ref, subgrid_size, image_size, w_step, shift.data());
+  reference.set_grid(grid_ref);
   reference.gridding(plan, w_offset, shift, cell_size, kernel_size,
                      subgrid_size, frequencies, visibilities, uvw, baselines,
                      aterms, aterms_offsets, spheroidal);
@@ -217,27 +216,25 @@ int compare_to_reference(float tol = 1000 *
                                   grid->data(), grid_ref->data());
 
   // Use the same grid for both degridding calls
-  reference.set_grid(optimized.get_grid(), subgrid_size, image_size, w_step,
-                     shift.data());
+  reference.set_grid(optimized.get_grid());
 #else
-  optimized.set_grid(grid, subgrid_size, image_size, w_step, shift.data());
-  reference.set_grid(grid, subgrid_size, image_size, w_step, shift.data());
+  optimized.set_grid(grid);
+  reference.set_grid(grid);
 #endif
-  reference.set_grid(optimized.get_grid(), subgrid_size, image_size, w_step,
-                     shift.data());
+  reference.set_grid(optimized.get_grid());
 
 #if TEST_DEGRIDDING
   // Run degridder
   std::clog << ">>> Run degridding" << std::endl;
   visibilities.zero();
   visibilities_ref.zero();
-  optimized.set_grid(grid, subgrid_size, image_size, w_step, shift.data());
+  optimized.set_grid(grid);
   optimized.degridding(plan, w_offset, shift, cell_size, kernel_size,
                        subgrid_size, frequencies, visibilities, uvw, baselines,
                        aterms, aterms_offsets, spheroidal);
 
   std::clog << ">>> Run reference degridding" << std::endl;
-  reference.set_grid(grid, subgrid_size, image_size, w_step, shift.data());
+  reference.set_grid(grid);
   reference.degridding(plan, w_offset, shift, cell_size, kernel_size,
                        subgrid_size, frequencies, visibilities_ref, uvw,
                        baselines, aterms, aterms_offsets, spheroidal);

idg-lib/src/CPU/common/CPU.cpp

@@ -65,20 +65,14 @@ std::unique_ptr<Plan> CPU::make_plan(
   }
 }
 
-void CPU::set_grid(std::shared_ptr<Grid> grid) {
-  Proxy::set_grid(grid);
-  m_wtiles = WTiles();
-}
-
-void CPU::set_grid(std::shared_ptr<Grid> grid, int subgrid_size,
-                   float image_size, float w_step, const float *shift) {
-  Proxy::set_grid(grid, subgrid_size, image_size, w_step, shift);
+void CPU::init_wtiles(float subgrid_size)
+{
   m_wtiles = WTiles(kernel::cpu::InstanceCPU::kNrWTiles,
                     kernel::cpu::InstanceCPU::kWTileSize);
   kernels.init_wtiles(subgrid_size);
 }
 
-std::shared_ptr<Grid> CPU::get_grid() {
+void CPU::flush_wtiles() {
   // flush all pending Wtiles
   WTileUpdateInfo wtile_flush_info = m_wtiles.clear();
   if (wtile_flush_info.wtile_ids.size()) {
@@ -87,7 +81,6 @@ std::shared_ptr<Grid> CPU::get_grid() {
         wtile_flush_info.wtile_ids.size(), wtile_flush_info.wtile_ids.data(),
         wtile_flush_info.wtile_coordinates.data(), m_grid->data());
   }
-  return m_grid;
 }
 
 unsigned int CPU::compute_jobsize(const Plan &plan,

idg-lib/src/CPU/common/CPU.h

@@ -47,13 +47,9 @@ class CPU : public Proxy {
       const Array1D<unsigned int>& aterms_offsets,
       Plan::Options options) override;
 
-  using Proxy::set_grid;  // prevents hiding set_grid overloads in Proxy
-  virtual void set_grid(std::shared_ptr<Grid> grid);
+  virtual void init_wtiles(float subgrid_size) override;
 
-  virtual void set_grid(std::shared_ptr<Grid> grid, int subgrid_size,
-                        float image_size, float w_step,
-                        const float* shift) override;
-  virtual std::shared_ptr<Grid> get_grid() override;
+  virtual void flush_wtiles() override;
 
  private:
   unsigned int compute_jobsize(const Plan& plan,

idg-lib/src/CUDA/Generic/Generic.cpp

@@ -478,12 +478,6 @@ void Generic::do_degridding(
 }  // end degridding
 
 void Generic::set_grid(std::shared_ptr<Grid> grid) {
-  set_grid(grid, 0, 0.0, 0.0, nullptr);
-}
-
-void Generic::set_grid(std::shared_ptr<Grid> grid, int /* subgrid_size */,
-                       float /* image_size */, float /* w_step */,
-                       const float* /* shift */) {
   m_grid = grid;
   InstanceCUDA& device = get_device(0);
   device.allocate_device_grid(grid->bytes());

idg-lib/src/CUDA/Generic/Generic.h

@@ -74,10 +74,7 @@ class Generic : public CUDA {
       const Array2D<float>& spheroidal);
 
  public:
-  virtual void set_grid(std::shared_ptr<Grid> grid);
-  virtual void set_grid(std::shared_ptr<Grid> grid, int subgrid_size,
-                        float image_size, float w_step,
-                        const float* shift) override;
+  virtual void set_grid(std::shared_ptr<Grid> grid) override;
 
   virtual std::shared_ptr<Grid> get_grid() override;
 

idg-lib/src/CUDA/Unified/Unified.cpp

@@ -106,12 +106,6 @@ void Unified::do_degridding(
 }  // end degridding
 
 void Unified::set_grid(std::shared_ptr<Grid> grid) {
-  set_grid(grid, 0, 0.0, 0.0, nullptr);
-}
-
-void Unified::set_grid(std::shared_ptr<Grid> grid, int /* subgrid_size */,
-                       float /* image_size */, float /* w_step */,
-                       const float* /* shift */) {
   m_grid = grid;
 
   if (m_enable_tiling) {

idg-lib/src/CUDA/Unified/Unified.h

@@ -54,10 +54,6 @@ class Unified : public Generic {
 
   virtual void set_grid(std::shared_ptr<Grid> grid) override;
 
-  virtual void set_grid(std::shared_ptr<Grid> grid, int subgrid_size,
-                        float image_size, float w_step,
-                        const float* shift) override;
-
   virtual std::shared_ptr<Grid> get_grid() override;
 
  private:

idg-lib/src/Hybrid/CUDA/GenericOptimized/GenericOptimized.cpp

@@ -1064,24 +1064,20 @@ std::unique_ptr<Plan> GenericOptimized::make_plan(
 }
 
 void GenericOptimized::set_grid(std::shared_ptr<Grid> grid) {
-  // Set grid both for CUDA proxy and CPU Proxy
+  // Defer call to cpuProxy
   cpuProxy->set_grid(grid);
   CUDA::set_grid(grid);
 }
 
-void GenericOptimized::set_grid(std::shared_ptr<Grid> grid, int subgrid_size,
-                                float image_size, float w_step,
-                                const float* shift) {
-  // Set grid both for CUDA proxy and CPU Proxy
-  // cpuProxy manages the wtiles state
-  cpuProxy->set_grid(grid, subgrid_size, image_size, w_step, shift);
-  CUDA::set_grid(grid, subgrid_size, image_size, w_step, shift);
+std::shared_ptr<Grid> GenericOptimized::get_grid() {
+  // Defer call to cpuProxy
+  return cpuProxy->get_grid();
 }
 
-std::shared_ptr<Grid> GenericOptimized::get_grid() {
+void GenericOptimized::init_wtiles(float subgrid_size) {
   // Defer call to cpuProxy
   // cpuProxy manages the wtiles state
-  return cpuProxy->get_grid();
+  cpuProxy->init_wtiles(subgrid_size);
 }
 
 }  // namespace hybrid

idg-lib/src/Hybrid/CUDA/GenericOptimized/GenericOptimized.h

@@ -31,13 +31,12 @@ class GenericOptimized : public cuda::CUDA {
 
   virtual bool supports_avg_aterm_correction() { return true; }
 
-  virtual void set_grid(std::shared_ptr<Grid> grid);
+  virtual void set_grid(std::shared_ptr<Grid> grid) override;
 
-  virtual void set_grid(std::shared_ptr<Grid> grid, int subgrid_size,
-                        float image_size, float w_step,
-                        const float* shift) override;
   virtual std::shared_ptr<Grid> get_grid() override;
 
+  virtual void init_wtiles(float subgrid_size) override;
+
  private:
   virtual void do_gridding(
       const Plan& plan,

idg-lib/src/common/Proxy.cpp

@@ -194,10 +194,9 @@ void Proxy::calibrate_init(
   //    subgrid_size, frequencies, visibilities, uvw, baselines,
   //    grid, aterms, aterms_offsets, spheroidal);
 
-  float image_size = grid.get_x_dim() * cell_size;
   auto grid_shared_ptr =
       std::shared_ptr<Grid>(const_cast<Grid*>(&grid), [](Grid*) {});
-  set_grid(grid_shared_ptr, subgrid_size, image_size, w_step, shift.data());
+  set_grid(grid_shared_ptr);
 
   int nr_w_layers;
 
@@ -533,16 +532,13 @@ std::shared_ptr<Grid> Proxy::allocate_grid(size_t nr_w_layers,
       new Grid(nr_w_layers, nr_correlations, height, width));
 }
 
-void Proxy::set_grid(std::shared_ptr<Grid> grid, int subgrid_size,
-                     float image_size, float w_step, const float* shift) {
-  m_grid = grid;
-  m_grid_size = m_grid->get_y_dim();
-  m_subgrid_size = subgrid_size;
-  m_image_size = image_size;
-  m_w_step = w_step;
-  m_shift[0] = shift[0];
-  m_shift[1] = shift[1];
-  m_shift[2] = shift[2];
+void Proxy::set_grid(std::shared_ptr<idg::Grid> grid) {
+  // Don't create a new shared_ptr when the grid data pointer is
+  // the same. This can be the case when the C-interface is used.
+  if (!m_grid || m_grid->data() != grid->data())
+  {
+    m_grid = grid;
+  }
 }
 
 std::shared_ptr<Grid> Proxy::get_grid() { return m_grid; }

idg-lib/src/common/Proxy.h

@@ -227,13 +227,14 @@ class Proxy {
                                               size_t nr_correlations,
                                               size_t height, size_t width);
 
-  virtual void set_grid(std::shared_ptr<Grid> grid) { m_grid = grid; }
-
-  virtual void set_grid(std::shared_ptr<Grid> grid, int subgrid_size,
-                        float image_size, float w_step, const float* shift);
+  virtual void set_grid(std::shared_ptr<Grid> grid);
 
   virtual std::shared_ptr<Grid> get_grid();
 
+  virtual void init_wtiles(float subgrid_size) {};
+
+  virtual void flush_wtiles() {};
+
   //! Method W-tiling
   virtual std::unique_ptr<Plan> make_plan(
       const int kernel_size, const int subgrid_size, const int grid_size,

idg-lib/src/common/ProxyC.cpp

@@ -160,17 +160,13 @@ void Proxy_calibrate_hessian_vector_product2(
           antenna_nr, aterms_, aterm_derivatives_, parameter_vector_);
 }
 
-void Proxy_transform(Proxy* p, int direction, std::complex<float>* grid,
-                     unsigned int grid_nr_correlations,
-                     unsigned int grid_height, unsigned int grid_width) {
+void Proxy_transform(Proxy* p, int direction) {
   if (direction != 0) {
     reinterpret_cast<idg::proxy::Proxy*>(p)->transform(
-        idg::ImageDomainToFourierDomain, grid, grid_nr_correlations,
-        grid_height, grid_width);
+        idg::ImageDomainToFourierDomain);
   } else {
     reinterpret_cast<idg::proxy::Proxy*>(p)->transform(
-        idg::FourierDomainToImageDomain, grid, grid_nr_correlations,
-        grid_height, grid_width);
+        idg::FourierDomainToImageDomain);
   }
 }
 
@@ -179,24 +175,38 @@ void Proxy_destroy(Proxy* p) { delete reinterpret_cast<idg::proxy::Proxy*>(p); }
 void* Proxy_allocate_grid(Proxy* p, unsigned int nr_correlations,
                           unsigned int grid_size) {
   const unsigned int nr_w_layers = 1;
-  auto grid_ptr = reinterpret_cast<idg::proxy::Proxy*>(p)->allocate_grid(
+  auto grid = reinterpret_cast<idg::proxy::Proxy*>(p)->allocate_grid(
       nr_w_layers, nr_correlations, grid_size, grid_size);
-
-  reinterpret_cast<idg::proxy::Proxy*>(p)->set_grid(grid_ptr);
-  return grid_ptr->data();
+  reinterpret_cast<idg::proxy::Proxy*>(p)->set_grid(grid);
+  return grid->data();
 }
 
-// TODO expose in header file, change arguments to ctypes only
-
-// void Proxy_set_grid(Proxy* p, idg::Grid& grid, subgrid_size, image_size,
-// w_step, shift) {
-//   std::shared_ptr<idg::Grid> grid_ptr(&grid);
-//   reinterpret_cast<idg::proxy::Proxy*>(p)->set_grid(grid_ptr, subgrid_size,
-//   image_size, w_step, shift);
-// }
+void Proxy_set_grid(Proxy* p,
+                    std::complex<float>* grid_ptr,
+                    unsigned int nr_w_layers,
+                    unsigned int nr_correlations,
+                    unsigned int height,
+                    unsigned int width)
+{
+    std::shared_ptr<idg::Grid> grid =
+        std::shared_ptr<idg::Grid>(new idg::Grid(grid_ptr, nr_w_layers, nr_correlations, height, width));
+    reinterpret_cast<idg::proxy::Proxy*>(p)->set_grid(grid);
+}
 
-// void Proxy_get_grid(Proxy* p, void* ptr) {
-//   auto grid = reinterpret_cast<idg::proxy::Proxy*>(p)->get_grid();
-//   memcpy(ptr, grid->data(), grid->bytes());
-// }
+void Proxy_get_grid(Proxy* p,
+                    std::complex<float>* grid_ptr,
+                    unsigned int nr_w_layers,
+                    unsigned int nr_correlations,
+                    unsigned int height,
+                    unsigned int width)
+{
+    std::shared_ptr<idg::Grid> grid =
+        reinterpret_cast<idg::proxy::Proxy*>(p)->get_grid();
+    assert(grid->get_w_dim() == nr_w_layers);
+    assert(grid->get_z_dim() == nr_correlations);
+    assert(grid->get_y_dim() == height);
+    assert(grid->get_x_dim() == width);
+    memcpy(grid_ptr, grid->data(), grid->bytes());
 }
+
+}  // end extern "C"

idg-lib/src/common/ProxyC.h

@@ -78,9 +78,7 @@ void Proxy_calibrate_update_hessian_vector_product2(
     const unsigned int nr_terms, std::complex<float>* aterms,
     std::complex<float>* aterm_derivatives, float* parameter_vector);
 
-void Proxy_transform(Proxy* p, int direction, std::complex<float>* grid,
-                     unsigned int grid_nr_correlations,
-                     unsigned int grid_height, unsigned int grid_width);
+void Proxy_transform(Proxy* p, int direction);
 
 void Proxy_destroy(Proxy* p);
 

idg-lib/src/python/Proxy.py

@@ -24,7 +24,6 @@ class Proxy(object):
         visibilities,
         uvw,
         baselines,
-        grid,
         aterms,
         aterms_offsets,
         spheroidal):
@@ -43,9 +42,6 @@ class Proxy(object):
         :param baselines: numpy.ndarray(
                 shape=(nr_baselines),
                 dtype=idg.baselinetype)
-        :param grid: numpy.ndarray(
-                shape=(nr_correlations, height, width),
-                dtype = idg.gridtype)
         :param aterms: numpy.ndarray(
                 shape=(nr_timeslots, nr_stations, height, width, nr_correlations),
                 dtype = idg.atermtype)
@@ -67,9 +63,6 @@ class Proxy(object):
         uvw_nr_coordinates           = 3
         baselines_nr_baselines       = baselines.shape[0]
         baselines_two                = 2
-        grid_nr_correlations         = grid.shape[0]
-        grid_height                  = grid.shape[1]
-        grid_width                   = grid.shape[2]
         aterms_nr_timeslots          = aterms.shape[0]
         aterms_nr_stations           = aterms.shape[1]
         aterms_aterm_height          = aterms.shape[2]
@@ -101,10 +94,6 @@ class Proxy(object):
             ctypes.c_void_p, # baselines
             ctypes.c_int,
             ctypes.c_int,
-            ctypes.c_void_p, # grid
-            ctypes.c_int,
-            ctypes.c_int,
-            ctypes.c_int,
             ctypes.c_void_p, # aterms
             ctypes.c_int,
             ctypes.c_int,
@@ -137,10 +126,6 @@ class Proxy(object):
             baselines.ctypes.data_as(ctypes.c_void_p),
             ctypes.c_int(baselines_nr_baselines),
             ctypes.c_int(baselines_two),
-            grid.ctypes.data_as(ctypes.c_void_p),
-            ctypes.c_int(grid_nr_correlations),
-            ctypes.c_int(grid_height),
-            ctypes.c_int(grid_width),
             aterms.ctypes.data_as(ctypes.c_void_p),
             ctypes.c_int(aterms_nr_timeslots),
             ctypes.c_int(aterms_nr_stations),
@@ -164,7 +149,6 @@ class Proxy(object):
         visibilities,
         uvw,
         baselines,
-        grid,
         aterms,
         aterms_offsets,
         spheroidal):
@@ -183,9 +167,6 @@ class Proxy(object):
         :param baselines: numpy.ndarray(
                 shape=(nr_baselines),
                 dtype=idg.baselinetype)
-        :param grid: numpy.ndarray(
-                shape=(nr_correlations, height, width),
-                dtype = idg.gridtype)
         :param aterms: numpy.ndarray(
                 shape=(nr_timeslots, nr_stations, height, width, nr_correlations),
                 dtype = idg.atermtype)
@@ -207,9 +188,6 @@ class Proxy(object):
         uvw_nr_coordinates           = 3
         baselines_nr_baselines       = baselines.shape[0]
         baselines_two                = 2
-        grid_nr_correlations         = grid.shape[0]
-        grid_height                  = grid.shape[1]
-        grid_width                   = grid.shape[2]
         aterms_nr_timeslots          = aterms.shape[0]
         aterms_nr_stations           = aterms.shape[1]
         aterms_aterm_height          = aterms.shape[2]
@@ -241,10 +219,6 @@ class Proxy(object):
             ctypes.c_void_p, # baselines
             ctypes.c_int,
             ctypes.c_int,
-            ctypes.c_void_p, # grid
-            ctypes.c_int,
-            ctypes.c_int,
-            ctypes.c_int,
             ctypes.c_void_p, # aterms
             ctypes.c_int,
             ctypes.c_int,
@@ -277,10 +251,6 @@ class Proxy(object):
             baselines.ctypes.data_as(ctypes.c_void_p),
             ctypes.c_int(baselines_nr_baselines),
             ctypes.c_int(baselines_two),
-            grid.ctypes.data_as(ctypes.c_void_p),
-            ctypes.c_int(grid_nr_correlations),
-            ctypes.c_int(grid_height),
-            ctypes.c_int(grid_width),
             aterms.ctypes.data_as(ctypes.c_void_p),
             ctypes.c_int(aterms_nr_timeslots),
             ctypes.c_int(aterms_nr_stations),
@@ -567,35 +537,20 @@ class Proxy(object):
 
     def transform(
         self,
-        direction,
-        grid):
+        direction):
         """
         Transform Fourier Domain<->Image Domain.
 
         :param direction: idg.FourierDomainToImageDomain or idg.ImageDomainToFourierDomain
-        :param grid: numpy.ndarray(
-                shape=(nr_correlations, height, width),
-                dtype = idg.gridtype)
         """
-        # extract dimesions
-        nr_correlations = grid.shape[0]
-        height          = grid.shape[1]
-        width           = grid.shape[2]
 
         # call C function to do the work
         self.lib.Proxy_transform.argtypes = [
                 ctypes.c_void_p,
-                ctypes.c_int,
-                ctypes.c_void_p,
-                ctypes.c_int,
                 ctypes.c_int]
         self.lib.Proxy_transform(
-            self.obj,
-            ctypes.c_int(direction),
-            grid.ctypes.data_as(ctypes.c_void_p),
-            ctypes.c_int(nr_correlations),
-            ctypes.c_int(height),
-            ctypes.c_int(width))
+            ctypes.c_void_p(self.obj),
+            ctypes.c_int(direction))
 
     def allocate_grid(
         self,
@@ -624,3 +579,79 @@ class Proxy(object):
 
         # Return grid
         return grid
+
+    def set_grid(
+        self,
+        grid):
+
+        """
+        Set grid to use in proxy
+
+        :param grid: numpy.ndarray(
+                shape=(nr_correlations, height, width),
+                dtype = idg.gridtype)
+        """
+
+        # Get dimensions
+        shape = grid.shape
+        nr_w_layers = 1
+        nr_correlations = shape[0]
+        grid_size = shape[1]
+        height = grid_size
+        width = grid_size
+
+        # Set argument types
+        self.lib.Proxy_set_grid.argtypes = [
+                ctypes.c_void_p,
+                ctypes.c_void_p,
+                ctypes.c_int,