Make the ghost points explicit in prolongation

3 files changed, 25 insertions, 14 deletions

diff --git a/ell_grid_solve.c b/ell_grid_solve.c
index 13d9072..7dcd41e 100644
--- a/ell_grid_solve.c
+++ b/ell_grid_solve.c
@@ -73,10 +73,10 @@ typedef struct EGSExactInternal {
 } EGSExactInternal;
 
 struct EGSInternal {
-    NDArray *u_base;
     NDArray *rhs_base;
     NDArray *residual_base;
 
+    NDArray *u_next_base;
     NDArray *u_next;
     int      u_next_valid;
 
@@ -331,9 +331,12 @@ static int solve_relax_step(EGSContext *ctx, int export_res)
     int u_next_valid = priv->u_next_valid;
 
     if (u_next_valid) {
-        NDArray *tmp = ctx->u;
+        NDArray *tmp      = ctx->u;
+        NDArray *tmp_base = ctx->u_base;
         ctx->u       = priv->u_next;
+        ctx->u_base  = priv->u_next_base;
         priv->u_next = tmp;
+        priv->u_next_base = tmp_base;
         priv->u_next_valid = 0;
     }
 
@@ -930,26 +933,26 @@ static int arrays_alloc(EGSContext *ctx, const size_t domain_size[2])
         size_padded[0] * MG2D_DIFF_COEFF_NB,
         size_padded[1],
     };
-    const size_t size_u[2] = {
-        size_padded[0] * 2,
-        size_padded[1],
-    };
     const Slice slice[2] = { SLICE(ghosts, -ghosts, 1),
                              SLICE(ghosts, -ghosts, 1) };
     int ret;
 
-    ret = mg2di_ndarray_alloc(&priv->u_base, 2, size_u, NDARRAY_ALLOC_ZERO);
+    ret = mg2di_ndarray_alloc(&ctx->u_base, 2, size_padded, NDARRAY_ALLOC_ZERO);
     if (ret < 0)
         return ret;
 
-    ret = mg2di_ndarray_slice(&ctx->u, priv->u_base,
+    ret = mg2di_ndarray_slice(&ctx->u, ctx->u_base,
                               (Slice [2]){ SLICE(ghosts, size_padded[0] - ghosts, 1),
                                            SLICE(ghosts, -ghosts, 1) });
     if (ret < 0)
         return ret;
 
-    ret = mg2di_ndarray_slice(&priv->u_next, priv->u_base,
-                              (Slice [2]){ SLICE(size_padded[0] + ghosts, size_u[0] - ghosts, 1),
+    ret = mg2di_ndarray_alloc(&priv->u_next_base, 2, size_padded, NDARRAY_ALLOC_ZERO);
+    if (ret < 0)
+        return ret;
+
+    ret = mg2di_ndarray_slice(&priv->u_next, priv->u_next_base,
+                              (Slice [2]){ SLICE(ghosts, size_padded[0] - ghosts, 1),
                                            SLICE(ghosts, -ghosts, 1) });
     if (ret < 0)
         return ret;
@@ -1079,8 +1082,9 @@ void mg2di_egs_free(EGSContext **pctx)
     exact_arrays_free(ctx);
 
     mg2di_ndarray_free(&ctx->u);
+    mg2di_ndarray_free(&ctx->u_base);
     mg2di_ndarray_free(&ctx->priv->u_next);
-    mg2di_ndarray_free(&ctx->priv->u_base);
+    mg2di_ndarray_free(&ctx->priv->u_next_base);
 
     mg2di_ndarray_free(&ctx->rhs);
     mg2di_ndarray_free(&ctx->priv->rhs_base);
diff --git a/ell_grid_solve.h b/ell_grid_solve.h
index 601eac9..7ce61a0 100644
--- a/ell_grid_solve.h
+++ b/ell_grid_solve.h
@@ -156,6 +156,11 @@ typedef struct EGSContext {
     NDArray *u;
 
     /**
+     * u including the ghost zones.
+     */
+    NDArray *u_base;
+
+    /**
      * Values of the right-hand side.
      *
      * Allocated by the solver in mg2di_egs_alloc(), owned by the solver.
diff --git a/mg2d.c b/mg2d.c
index 84985d7..25a1759 100644
--- a/mg2d.c
+++ b/mg2d.c
@@ -192,7 +192,7 @@ static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level, int allow_exact)
             /* prolongate the coarser-level correction */
             mg2di_timer_start(&level->timer_prolong);
             ret = mg2di_gt_transfer(level->transfer_prolong, level->prolong_tmp,
-                                    level->child->solver->u);
+                                    level->child->solver->u_base);
             mg2di_timer_stop(&level->timer_prolong);
             if (ret < 0)
                 return ret;
@@ -444,8 +444,10 @@ static int mg_levels_init(MG2DContext *ctx)
             cur->transfer_prolong->dst.step[0] = cur->solver->step[0];
             cur->transfer_prolong->dst.step[1] = cur->solver->step[1];
 
-            cur->transfer_prolong->src.size[0] = cur->child->solver->domain_size[0];
-            cur->transfer_prolong->src.size[1] = cur->child->solver->domain_size[1];
+            cur->transfer_prolong->src.start[0] = -FD_STENCIL_MAX;
+            cur->transfer_prolong->src.start[1] = -FD_STENCIL_MAX;
+            cur->transfer_prolong->src.size[0] = cur->child->solver->domain_size[0] + 2 * FD_STENCIL_MAX;
+            cur->transfer_prolong->src.size[1] = cur->child->solver->domain_size[1] + 2 * FD_STENCIL_MAX;
             cur->transfer_prolong->src.step[0] = cur->child->solver->step[0];
             cur->transfer_prolong->src.step[1] = cur->child->solver->step[1];