mg2d: add API for specifying singular diff coeffs at the boundaries

API and ABI break.

4 files changed, 394 insertions, 91 deletions

diff --git a/mg2d.c b/mg2d.c
index e0002af..9775c16 100644
--- a/mg2d.c
+++ b/mg2d.c
@@ -109,7 +109,14 @@ struct MG2DInternal {
 
     NDArray *u;
     NDArray *rhs;
-    NDArray *diff_coeffs[MG2D_DIFF_COEFF_NB];
+    NDArray *diff_coeffs_base[MG2D_DIFF_COEFF_NB];
+    NDArray *diff_coeffs_tmp[MG2D_DIFF_COEFF_NB];
+
+    /**
+     * This is the full boundary val spanning all the components.
+     * A pointer inside it is exported as MG2DDCBoundarySpec.val
+     */
+    NDArray *dc_bnd_val[MG2D_DIFF_COEFF_NB][4];
 
     GridTransferContext *transfer_init;
 
@@ -356,6 +363,47 @@ static void bnd_copy(MG2DBoundary *bdst, const double *src, ptrdiff_t src_stride
     }
 }
 
+static void mirror(double *dst, const ptrdiff_t dst_stride[2],
+                   const size_t size[2], double parity)
+{
+    if (dst_stride[1] == 1 && parity == 1.0) {
+        for (int j = 1; j <= size[0]; j++)
+            memcpy(dst + j * dst_stride[0], dst - j * dst_stride[0], sizeof(*dst) * size[1]);
+    } else {
+        for (size_t i = 0; i < size[1]; i++) {
+            for (int j = 1; j <= size[0]; j++)
+                dst[dst_stride[1] * i + dst_stride[0] * j] = parity * dst[dst_stride[1] * i - dst_stride[0] * j];
+        }
+    }
+}
+
+static void dc_boundaries_apply(NDArray *a, const int *reflect, const double *parity,
+                                const DomainComponent *dc)
+{
+    for (int order = 0; order < 2; order++) {
+        for (int bnd_idx = 0; bnd_idx < 4; bnd_idx++) {
+            const int ci = mg2d_bnd_coord_idx(bnd_idx);
+            double  *dst = a->data + mg2d_bnd_is_upper(bnd_idx) * ((a->shape[!ci] - 1) * a->stride[!ci]);
+            const ptrdiff_t dst_strides[2] = { mg2d_bnd_out_dir(bnd_idx) * a->stride[!ci], a->stride[ci] };
+
+            if (!dc->bnd_is_outer[bnd_idx])
+                continue;
+
+            if (order == 0 && !reflect[bnd_idx]) {
+
+                for (ptrdiff_t bnd_layer = 1; bnd_layer <= FD_STENCIL_MAX; bnd_layer++)
+                    for (ptrdiff_t row_idx = -(ptrdiff_t)FD_STENCIL_MAX; row_idx < (ptrdiff_t)a->shape[ci] + FD_STENCIL_MAX; row_idx++) {
+                        dst[bnd_layer * dst_strides[0] + row_idx * dst_strides[1]] = 2.0 * dst[(bnd_layer - 1) * dst_strides[0] + row_idx * dst_strides[1]] -
+                                                                                           dst[(bnd_layer - 2) * dst_strides[0] + row_idx * dst_strides[1]];
+                    }
+            } else if (order == 1 && reflect[bnd_idx]) {
+                const size_t size[2] = { FD_STENCIL_MAX, a->shape[ci] + 2 * FD_STENCIL_MAX };
+                mirror(dst - FD_STENCIL_MAX * dst_strides[1], dst_strides, size, parity[bnd_idx]);
+            }
+        }
+    }
+}
+
 static int restrict_dc_sync(MG2DContext *ctx, MG2DLevel *l)
 {
     MG2DInternal *priv = ctx->priv;
@@ -364,7 +412,8 @@ static int restrict_dc_sync(MG2DContext *ctx, MG2DLevel *l)
     int *senddispl = NULL, *recvdispl = NULL, *sendcounts = NULL, *recvcounts = NULL;
 
     const ptrdiff_t *lo = l->dg->components[priv->local_component].interior.start;
-    const ptrdiff_t stride = l->solver->diff_coeffs[0]->stride[0];
+    NDArray **diff_coeffs = l->depth > 0 ? l->solver->diff_coeffs : priv->diff_coeffs_tmp;
+    const ptrdiff_t stride = diff_coeffs[0]->stride[0];
 
     int ret = 0;
 
@@ -416,11 +465,9 @@ static int restrict_dc_sync(MG2DContext *ctx, MG2DLevel *l)
                          (overlaps_send[i].start[0] - lo[0])) * sizeof(double);
     }
 
-    for (int i = 0; i < ARRAY_ELEMS(l->solver->diff_coeffs); i++) {
-        MPI_Alltoallw(l->solver->diff_coeffs[i]->data,
-                      sendcounts, senddispl, sendtypes,
-                      l->solver->diff_coeffs[i]->data,
-                      recvcounts, recvdispl, recvtypes,
+    for (int i = 0; i < MG2D_DIFF_COEFF_NB; i++) {
+        MPI_Alltoallw(diff_coeffs[i]->data, sendcounts, senddispl, sendtypes,
+                      diff_coeffs[i]->data, recvcounts, recvdispl, recvtypes,
                       l->mpi_comm);
     }
 
@@ -448,50 +495,113 @@ fail:
     return ret;
 }
 
-static int restrict_diff_coeffs(MG2DContext *ctx, MG2DLevel *l,
-                                enum GridTransferOperator op)
+static int restrict_diff_coeffs(MG2DContext *ctx, MG2DLevel *l)
 {
     MG2DInternal            *priv = ctx->priv;
     const DomainComponent *dc_src = &l->dg->components[priv->local_component];
     int           multi_component = l->dg->nb_components > 1;
 
-    GridTransferContext *tc;
     int ret = 0;
+    int   reflect[MG2D_DIFF_COEFF_NB][4];
+    double parity[MG2D_DIFF_COEFF_NB][4];
 
-    /* get the ghost points for interpolation */
-    if (multi_component) {
-        ret = restrict_dc_sync(ctx, l);
-        if (ret < 0)
-            return ret;
-    }
+    /* prepare the source arrays */
+    for (int dc_idx = 0; dc_idx < ARRAY_ELEMS(ctx->diff_coeffs); dc_idx++) {
+        MG2DDiffCoeffs *dc = ctx->diff_coeffs[dc_idx];
+        NDArray *a_dc = l->depth > 0 ? l->solver->diff_coeffs[dc_idx] : priv->diff_coeffs_tmp[dc_idx];
 
-    tc = mg2di_gt_alloc(op);
-    if (!tc)
-        return -ENOMEM;
+        if (!l->depth)
+            mg2di_ndarray_copy(a_dc, priv->root->solver->diff_coeffs[dc_idx]);
 
-    tc->tp                   = priv->tp;
-    tc->cpuflags             = priv->cpuflags;
-    tc->extrapolate_distance = 1;
+        for (int bnd_idx = 0; bnd_idx < ARRAY_ELEMS(dc->boundaries); bnd_idx++) {
+            const int ci = mg2d_bnd_coord_idx(bnd_idx);
+            double  *dst = a_dc->data + mg2d_bnd_is_upper(bnd_idx) * ((a_dc->shape[!ci] - 1) * a_dc->stride[!ci]);
+
+            reflect[dc_idx][bnd_idx] = ctx->boundaries[bnd_idx]->type == MG2D_BC_TYPE_REFLECT;
+            parity [dc_idx][bnd_idx] = 1.0;
+
+            if (dc_idx == MG2D_DIFF_COEFF_11 ||
+                (ci == 0 && dc_idx == MG2D_DIFF_COEFF_10) ||
+                (ci == 1 && dc_idx == MG2D_DIFF_COEFF_01))
+                parity[dc_idx][bnd_idx] *= -1.0;
+            if (dc->boundaries[bnd_idx].flags & MG2D_DC_FLAG_POLE)
+                parity[dc_idx][bnd_idx] *= -1.0;
+
+            /* gather the boundary values for interpolation onto coarser levels */
+            if (!l->depth && multi_component &&
+                dc->boundaries[bnd_idx].flags &MG2D_DC_FLAG_DISCONT) {
+                int *recvcounts = NULL, *displs = NULL;
+
+                recvcounts = calloc(l->dg->nb_components, sizeof(*recvcounts));
+                displs     = calloc(l->dg->nb_components, sizeof(*displs));
+                if (!recvcounts || !displs) {
+                    free(recvcounts);
+                    free(displs);
+                    return -ENOMEM;
+                }
 
-    for (int dim = 0; dim < 2; dim++) {
-        tc->src.start[!dim] = dc_src->interior.start[dim];
-        tc->src.size[!dim]  = dc_src->interior.size[dim];
-        tc->src.step[!dim]  = priv->dh.stepsizes[l->depth][dim];
+                for (int comp = 0; comp < l->dg->nb_components; comp++) {
+                    if (l->dg->components[comp].bnd_is_outer[bnd_idx]) {
+                        recvcounts[comp] = l->dg->components[comp].interior.size[!ci];
+                        displs[comp]     = l->dg->components[comp].interior.start[!ci];
+                    } else {
+                        recvcounts[comp] = 0;
+                        displs[comp]     = 0;
+                    }
+                }
 
-        tc->dst.start[!dim] = l->restrict_dst_extents.start[dim];
-        tc->dst.size[!dim]  = l->restrict_dst_extents.size[dim];
-        tc->dst.step[!dim]  = priv->dh.stepsizes[l->depth + 1][dim];
+                MPI_Allgatherv(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL,
+                               priv->dc_bnd_val[dc_idx][bnd_idx]->data + FD_STENCIL_MAX,
+                               recvcounts, displs, MPI_DOUBLE, priv->mpi_comm);
+                free(recvcounts);
+                free(displs);
+            }
+
+            if (l->depth > 0)
+                continue;
+
+            // Make the boundaries with poles continuous by multiplying by x
+            // FIXME only implemented for 0L boundary
+            if (bnd_idx == MG2D_BOUNDARY_0L &&
+                dc->boundaries[bnd_idx].flags & MG2D_DC_FLAG_POLE) {
+                const ptrdiff_t x_offset = l->dg->components[priv->local_component].interior.start[0];
+                for (size_t idx0 = 0; idx0 < a_dc->shape[0]; idx0++) {
+                    if (!x_offset)
+                        a_dc->data[idx0 * a_dc->stride[0]] = dc->boundaries[MG2D_BOUNDARY_0L].val[idx0];
+                    for (size_t idx1 = !x_offset; idx1 < a_dc->shape[1]; idx1++) {
+                        const double x = ctx->step[0] * (x_offset + idx1);
+                        a_dc->data[idx0 * a_dc->stride[0] + idx1 * a_dc->stride[1]] *= x;
+                    }
+                }
+            }
+
+            // Make the boundaries with discontinuities continous by
+            // substracting the discontinous jump
+            if (dc->boundaries[bnd_idx].flags & MG2D_DC_FLAG_DISCONT &&
+                l->dg->components[priv->local_component].bnd_is_outer[bnd_idx]) {
+                const double *src = dc->boundaries[bnd_idx].val;
+                for (int j = 0; j < a_dc->shape[ci]; j++)
+                    dst[j * a_dc->stride[ci]] -= src[j];
+            }
+        }
+
+        /* fill the boundary ghost zones for interpolation */
+        dc_boundaries_apply(a_dc, reflect[dc_idx], parity[dc_idx], dc_src);
     }
 
-    ret = mg2di_gt_init(tc);
-    if (ret < 0)
-        goto finish;
+    /* get the ghost points for interpolation */
+    if (multi_component) {
+        ret = restrict_dc_sync(ctx, l);
+        if (ret < 0)
+            return ret;
+    }
 
+    /* execute the interpolations on the continuous data */
     for (int i = 0; i < MG2D_DIFF_COEFF_NB; i++) {
-        ret = mg2di_gt_transfer(tc, multi_component ? l->restrict_dst : l->child->solver->diff_coeffs[i],
-                                l->solver->diff_coeffs[i]);
+        ret = mg2di_gt_transfer(l->transfer_restrict, multi_component ? l->restrict_dst : l->child->solver->diff_coeffs[i],
+                                l->depth > 0 ? l->solver->diff_coeffs[i] : priv->diff_coeffs_tmp[i]);
         if (ret < 0)
-            goto finish;
+            return ret;
 
         if (multi_component) {
             mg2di_timer_start(&l->timer_mpi_sync);
@@ -504,11 +614,100 @@ static int restrict_diff_coeffs(MG2DContext *ctx, MG2DLevel *l,
         }
     }
 
+    return ret;
+}
 
-finish:
-    mg2di_gt_free(&tc);
+/* return the coeffs into their final form */
+static int diff_coeffs_fixup(MG2DContext *ctx, MG2DLevel *l)
+{
+    MG2DInternal *priv = ctx->priv;
+    int ret;
 
-    return ret;
+    for (int i = 0; i < ARRAY_ELEMS(l->solver->diff_coeffs); i++) {
+        MG2DDiffCoeffs *dc = ctx->diff_coeffs[i];
+        NDArray     *a_dst = l->solver->diff_coeffs[i];
+
+        for (int bnd_idx = 0; bnd_idx < ARRAY_ELEMS(dc->boundaries); bnd_idx++) {
+            const int ci = mg2d_bnd_coord_idx(bnd_idx);
+
+            if (!l->dg->components[priv->local_component].bnd_is_outer[bnd_idx])
+                continue;
+
+            if (dc->boundaries[bnd_idx].flags & MG2D_DC_FLAG_DISCONT) {
+                GridTransferContext *gt_bnd = NULL;
+                NDArray *dc_src = NULL, *dc_dst = NULL;
+
+                double  *dst = a_dst->data + mg2d_bnd_is_upper(bnd_idx) *
+                                             ((a_dst->shape[!ci] - 1) * a_dst->stride[!ci]);
+                const ptrdiff_t stride_dst = a_dst->stride[ci];
+                const size_t      size_dst[] = { 1, l->solver->domain_size[!ci] };
+
+                ret = mg2di_ndarray_slice(&dc_src, priv->dc_bnd_val[i][bnd_idx],
+                                          (Slice [2]){ SLICE_NULL, SLICE(FD_STENCIL_MAX, -FD_STENCIL_MAX, 1) });
+                if (ret < 0)
+                    goto fail;
+
+                ret = mg2di_ndarray_alloc(&dc_dst, 2, size_dst, 0);
+                if (ret < 0)
+                    goto fail;
+
+                gt_bnd = mg2di_gt_alloc(GRID_TRANSFER_LAGRANGE_1);
+                if (!gt_bnd) {
+                    ret = -ENOMEM;
+                    goto fail;
+                }
+
+                gt_bnd->tp       = priv->tp;
+                gt_bnd->cpuflags = priv->cpuflags;
+
+                gt_bnd->src.start[0] = 0;
+                gt_bnd->src.start[1] = 0;
+                gt_bnd->src.size[0] = 1;
+                gt_bnd->src.size[1] = priv->dg->domain_size[1];
+                gt_bnd->src.step[0] = ctx->step[0];
+                gt_bnd->src.step[1] = ctx->step[1];
+
+                gt_bnd->dst.start[0] = 0;
+                gt_bnd->dst.start[1] = l->dg->components[priv->local_component].interior.start[!ci];
+                gt_bnd->dst.size[0] = 1;
+                gt_bnd->dst.size[1] = l->dg->components[priv->local_component].interior.size[!ci];
+                gt_bnd->dst.step[0] = l->solver->step[0];
+                gt_bnd->dst.step[1] = l->solver->step[1];
+
+                ret = mg2di_gt_init(gt_bnd);
+                if (ret < 0)
+                    goto fail;
+
+                ret = mg2di_gt_transfer(gt_bnd, dc_dst, dc_src);
+                if (ret < 0)
+                    goto fail;
+
+                for (int idx = 0; idx < dc_dst->shape[1]; idx++)
+                    dst[stride_dst * idx] += dc_dst->data[idx];
+
+fail:
+                mg2di_gt_free(&gt_bnd);
+                mg2di_ndarray_free(&dc_src);
+                mg2di_ndarray_free(&dc_dst);
+                if (ret < 0)
+                    return ret;
+            }
+        }
+
+        if (dc->boundaries[MG2D_BOUNDARY_0L].flags & MG2D_DC_FLAG_POLE) {
+            const ptrdiff_t x_offset = l->dg->components[priv->local_component].interior.start[0];
+            for (size_t idx0 = 0; idx0 < a_dst->shape[0]; idx0++) {
+                if (!x_offset)
+                    a_dst->data[idx0 * a_dst->stride[0]] = 0.0;
+                for (size_t idx1 = !x_offset; idx1 < a_dst->shape[1]; idx1++) {
+                    const double x = l->solver->step[0] * (x_offset + idx1);
+                    a_dst->data[idx0 * a_dst->stride[0] + idx1 * a_dst->stride[1]] /= x;
+                }
+            }
+        }
+    }
+
+    return 0;
 }
 
 static double findmax(double *arr, size_t size[2], ptrdiff_t stride)
@@ -597,12 +796,12 @@ static int mg_levels_init(MG2DContext *ctx)
     double diff0_max, diff2_max, tmp;
     int ret;
 
-    diff0_max = findmax(ctx->diff_coeffs[MG2D_DIFF_COEFF_00], priv->root->solver->domain_size,
-                        ctx->diff_coeffs_stride);
-    diff2_max = findmax(ctx->diff_coeffs[MG2D_DIFF_COEFF_20], priv->root->solver->domain_size,
-                        ctx->diff_coeffs_stride);
-    tmp = findmax(ctx->diff_coeffs[MG2D_DIFF_COEFF_02], priv->root->solver->domain_size,
-                              ctx->diff_coeffs_stride);
+    diff0_max = findmax(ctx->diff_coeffs[MG2D_DIFF_COEFF_00]->data, priv->root->solver->domain_size,
+                        ctx->diff_coeffs[MG2D_DIFF_COEFF_00]->stride);
+    diff2_max = findmax(ctx->diff_coeffs[MG2D_DIFF_COEFF_20]->data, priv->root->solver->domain_size,
+                        ctx->diff_coeffs[MG2D_DIFF_COEFF_20]->stride);
+    tmp = findmax(ctx->diff_coeffs[MG2D_DIFF_COEFF_02]->data, priv->root->solver->domain_size,
+                  ctx->diff_coeffs[MG2D_DIFF_COEFF_02]->stride);
     diff2_max = MAX(diff2_max, tmp);
 
     if (priv->dg->nb_components > 1) {
@@ -633,13 +832,12 @@ static int mg_levels_init(MG2DContext *ctx)
         ctx->rhs_stride = cur->solver->rhs->stride[0];
     }
 
-    if (priv->diff_coeffs[0]) {
+    if (ctx->diff_coeffs[0]->data == priv->diff_coeffs_tmp[0]->data) {
         for (int i = 0; i < ARRAY_ELEMS(ctx->diff_coeffs); i++) {
-            mg2di_ndarray_copy(cur->solver->diff_coeffs[i], priv->diff_coeffs[i]);
-            mg2di_ndarray_free(&priv->diff_coeffs[i]);
-            ctx->diff_coeffs[i] = cur->solver->diff_coeffs[i]->data;
+            mg2di_ndarray_copy(cur->solver->diff_coeffs[i], priv->diff_coeffs_tmp[i]);
+            ctx->diff_coeffs[i]->data   = cur->solver->diff_coeffs[i]->data;
+            ctx->diff_coeffs[i]->stride = cur->solver->diff_coeffs[i]->stride[0];
         }
-        ctx->diff_coeffs_stride = cur->solver->diff_coeffs[0]->stride[0];
     }
 
     while (cur) {
@@ -702,8 +900,9 @@ static int mg_levels_init(MG2DContext *ctx)
                     op_restrict = GRID_TRANSFER_FW_2;
                 else
                     op_restrict = GRID_TRANSFER_FW_3;
-            } else
+            } else {
                 op_restrict = op_interp;
+            }
 
             ret = mg_setup_restrict(ctx, cur, op_restrict);
             if (ret < 0)
@@ -713,13 +912,22 @@ static int mg_levels_init(MG2DContext *ctx)
             if (ret < 0)
                 return ret;
 
-            ret = restrict_diff_coeffs(ctx, cur, op_interp);
+            ret = restrict_diff_coeffs(ctx, cur);
             if (ret < 0)
                 return ret;
         }
 
-        cur->egs_init_flags &= ~EGS_INIT_FLAG_SAME_DIFF_COEFFS;
+        cur = cur->child;
+    }
 
+    cur = priv->root;
+    while (cur) {
+        if (cur->depth > 0) {
+            ret = diff_coeffs_fixup(ctx, cur);
+            if (ret < 0)
+                return ret;
+        }
+        cur->egs_init_flags &= ~EGS_INIT_FLAG_SAME_DIFF_COEFFS;
         cur = cur->child;
     }
 
@@ -1357,14 +1565,15 @@ static void log_default_callback(const MG2DContext *ctx, int level, const char *
     vfprintf(stderr, fmt, vl);
 }
 
-static MG2DContext *solver_alloc(const size_t domain_size[2])
+static MG2DContext *solver_alloc(DomainGeometry *dg, unsigned int local_component)
 {
+    const size_t *domain_size = dg->components[local_component].interior.size;
     MG2DContext *ctx;
     MG2DInternal *priv;
     int ret;
 
-    if (domain_size[0] < 3 || domain_size[1] < 3 ||
-        SIZE_MAX / domain_size[0] < domain_size[1])
+    if (dg->domain_size[0] < 3 || dg->domain_size[1] < 3 ||
+        SIZE_MAX / dg->domain_size[0] < dg->domain_size[1])
         return NULL;
 
     ctx = calloc(1, sizeof(*ctx));
@@ -1376,6 +1585,9 @@ static MG2DContext *solver_alloc(const size_t domain_size[2])
         goto fail;
     priv = ctx->priv;
 
+    priv->dg              = dg;
+    priv->local_component = local_component;
+
     priv->logger.log    = log_callback;
     priv->logger.opaque = ctx;
 
@@ -1404,13 +1616,42 @@ static MG2DContext *solver_alloc(const size_t domain_size[2])
     ctx->rhs_stride = priv->rhs->stride[0];
 
     for (int i = 0; i < ARRAY_ELEMS(ctx->diff_coeffs); i++) {
-        ret = mg2di_ndarray_alloc(&priv->diff_coeffs[i], 2, (size_t [2]){ domain_size[1], domain_size[0] },
-                                  NDARRAY_ALLOC_ZERO);
+        MG2DDiffCoeffs *dc;
+        const Slice slice[2] = { SLICE(FD_STENCIL_MAX, -FD_STENCIL_MAX, 1),
+                                 SLICE(FD_STENCIL_MAX, -FD_STENCIL_MAX, 1) };
+
+        ret = mg2di_ndarray_alloc(&priv->diff_coeffs_base[i], 2,
+                                  (size_t [2]){ domain_size[1] + 2 * FD_STENCIL_MAX,
+                                                domain_size[0] + 2 * FD_STENCIL_MAX },
+                                   NDARRAY_ALLOC_ZERO);
         if (ret < 0)
             goto fail;
-        ctx->diff_coeffs[i] = priv->diff_coeffs[i]->data;
+
+        ret = mg2di_ndarray_slice(&priv->diff_coeffs_tmp[i], priv->diff_coeffs_base[i], slice);
+        if (ret < 0)
+            goto fail;
+
+        dc = calloc(1, sizeof(*dc));
+        if (!dc)
+            goto fail;
+        ctx->diff_coeffs[i] = dc;
+
+        dc->data   = priv->diff_coeffs_tmp[i]->data;
+        dc->stride = priv->diff_coeffs_tmp[i]->stride[0];
+
+        for (int bnd_idx = 0; bnd_idx < ARRAY_ELEMS(dc->boundaries); bnd_idx++) {
+            const int ci = mg2d_bnd_coord_idx(bnd_idx);
+
+            ret = mg2di_ndarray_alloc(&priv->dc_bnd_val[i][bnd_idx], 2,
+                                      (size_t [2]){ 1, dg->domain_size[!ci] + 2 * FD_STENCIL_MAX },
+                                      NDARRAY_ALLOC_ZERO);
+            if (ret < 0)
+                goto fail;
+
+            dc->boundaries[bnd_idx].val = priv->dc_bnd_val[i][bnd_idx]->data + FD_STENCIL_MAX +
+                                          dg->components[local_component].interior.start[!ci];
+        }
     }
-    ctx->diff_coeffs_stride = priv->diff_coeffs[0]->stride[0];
 
     ctx->max_levels    = 16;
     ctx->max_exact_size = 5;
@@ -1435,32 +1676,32 @@ fail:
 MG2DContext *mg2d_solver_alloc(size_t domain_size)
 {
     MG2DContext *ctx;
+    DomainGeometry *dg;
 
-    ctx = solver_alloc((size_t [2]){ domain_size, domain_size });
-    if (!ctx)
+    dg = mg2di_dg_alloc(1);
+    if (!dg)
         return NULL;
 
-    ctx->priv->dg = mg2di_dg_alloc(1);
-    if (!ctx->priv->dg)
-        goto fail;
-
-    ctx->priv->dg->domain_size[0] = domain_size;
-    ctx->priv->dg->domain_size[1] = domain_size;
+    dg->domain_size[0] = domain_size;
+    dg->domain_size[1] = domain_size;
 
-    ctx->priv->dg->components[0].interior.start[0] = 0;
-    ctx->priv->dg->components[0].interior.start[1] = 0;
-    ctx->priv->dg->components[0].interior.size[0] = domain_size;
-    ctx->priv->dg->components[0].interior.size[1] = domain_size;
+    dg->components[0].interior.start[0] = 0;
+    dg->components[0].interior.start[1] = 0;
+    dg->components[0].interior.size[0] = domain_size;
+    dg->components[0].interior.size[1] = domain_size;
 
-    ctx->priv->dg->components[0].exterior.start[0] = -FD_STENCIL_MAX;
-    ctx->priv->dg->components[0].exterior.start[1] = -FD_STENCIL_MAX;
-    ctx->priv->dg->components[0].exterior.size[0]  = domain_size + 2 * FD_STENCIL_MAX;
-    ctx->priv->dg->components[0].exterior.size[1]  = domain_size + 2 * FD_STENCIL_MAX;
+    dg->components[0].exterior.start[0] = -FD_STENCIL_MAX;
+    dg->components[0].exterior.start[1] = -FD_STENCIL_MAX;
+    dg->components[0].exterior.size[0]  = domain_size + 2 * FD_STENCIL_MAX;
+    dg->components[0].exterior.size[1]  = domain_size + 2 * FD_STENCIL_MAX;
     for (int i = 0; i < 4; i++)
-        ctx->priv->dg->components[0].bnd_is_outer[i] = 1;
+        dg->components[0].bnd_is_outer[i] = 1;
 
-    ctx->priv->local_component = 0;
-    ctx->priv->mpi_comm        = MPI_COMM_NULL;
+    ctx = solver_alloc(dg, 0);
+    if (!ctx)
+        goto fail;
+
+    ctx->priv->mpi_comm = MPI_COMM_NULL;
 
     ctx->domain_size = domain_size;
 
@@ -1472,7 +1713,7 @@ MG2DContext *mg2d_solver_alloc(size_t domain_size)
 
     return ctx;
 fail:
-    mg2d_solver_free(&ctx);
+    mg2di_dg_free(&dg);
     return NULL;
 }
 
@@ -1541,13 +1782,11 @@ MG2DContext *mg2d_solver_alloc_mpi(MPI_Comm comm, const size_t local_start[2],
 
     free(domainspec);
 
-    ctx = solver_alloc(local_size);
+    ctx = solver_alloc(dg, rank);
     if (!ctx)
         return NULL;
 
     ctx->priv->mpi_comm        = comm;
-    ctx->priv->dg              = dg;
-    ctx->priv->local_component = rank;
     ctx->domain_size           = dg->domain_size[0];
     ctx->local_start[0]        = local_start[0];
     ctx->local_start[1]        = local_start[1];
@@ -1584,8 +1823,14 @@ void mg2d_solver_free(MG2DContext **pctx)
 
     mg2di_ndarray_free(&ctx->priv->u);
     mg2di_ndarray_free(&ctx->priv->rhs);
-    for (int i = 0; i < ARRAY_ELEMS(ctx->priv->diff_coeffs); i++)
-        mg2di_ndarray_free(&ctx->priv->diff_coeffs[i]);
+    for (int i = 0; i < ARRAY_ELEMS(ctx->priv->diff_coeffs_base); i++) {
+        for (int j = 0; j < ARRAY_ELEMS(ctx->priv->dc_bnd_val[i]); j++)
+            mg2di_ndarray_free(&ctx->priv->dc_bnd_val[i][j]);
+        free(ctx->diff_coeffs[i]);
+
+        mg2di_ndarray_free(&ctx->priv->diff_coeffs_tmp[i]);
+        mg2di_ndarray_free(&ctx->priv->diff_coeffs_base[i]);
+    }
 
     mg2di_dg_free(&ctx->priv->dg);
 
diff --git a/mg2d.h b/mg2d.h
index 47d709e..2f593bc 100644
--- a/mg2d.h
+++ b/mg2d.h
@@ -29,6 +29,68 @@
 
 typedef struct MG2DInternal MG2DInternal;
 
+/**
+ * Flags for diff coeffs boundary specifications.
+ */
+enum MG2DDCFlag {
+    /**
+     * The coefficients have a discontinuity at the boundary, i.e.
+     * C(x, y) = C0(x, y) + D(x, y),
+     * where C is the resulting value used for the solve, C0 is continuous
+     * everywhere, and D is zero everywhere except along the boundary. D is
+     * assumed to be continuous along the boundary.
+     *
+     * The values of D should be stored in MG2DDCBoundarySpec.val
+     */
+    MG2D_DC_FLAG_DISCONT = (1 << 0),
+    /**
+     * The coefficients have a first order pole at the boundary, i.e.
+     * C(x, y) · (x_i - x_i^0)                  (1)
+     * is continuous, where C is the actual value of the coefficient, x_i is the
+     * coordinate which is constant along the boundary and x_i^0 is the value of
+     * that coordinate along the boundary.
+     *
+     * The value of (1) along the boundary should be stored in
+     * MG2DDCBoundarySpec.val.
+     */
+    MG2D_DC_FLAG_POLE    = (1 << 1),
+};
+
+/**
+ * Specification of the behaviour of differential equation coefficients on a
+ * boundary.
+ */
+typedef struct MG2DDCBoundarySpec {
+    /**
+     * A combination of MG2DDCFlag
+     */
+    int flags;
+    /**
+     * Value determined by flags. Need not be filled if the flags are zero.
+     */
+    double *val;
+} MG2DDCBoundarySpec;
+
+/**
+ * Specification of a single coefficient of the PDE.
+ */
+typedef struct MG2DDiffCoeffs {
+    /**
+     * Values of the coefficient at the grid points. Values corresponding to
+     * successive x locations are contiguous in memory.
+     */
+    double         *data;
+    /**
+     * Number of elements between successive y locations in data.
+     */
+    ptrdiff_t       stride;
+
+    /**
+     * Behaviour of the coefficient on the boundaries.
+     */
+    MG2DDCBoundarySpec boundaries[4];
+} MG2DDiffCoeffs;
+
 typedef struct MG2DContext {
     /**
      * Solver private data, not to be accessed in any way by the caller.
@@ -127,11 +189,7 @@ typedef struct MG2DContext {
      * Allocated by the solver in mg2d_solver_alloc(), owned by the solver.
      * Must be filled by the caller before solving.
      */
-    double   *diff_coeffs[MG2D_DIFF_COEFF_NB];
-    /**
-     * Distance between neighbouring gridpoints along coord1.
-     */
-    ptrdiff_t diff_coeffs_stride;
+    MG2DDiffCoeffs *diff_coeffs[MG2D_DIFF_COEFF_NB];
 
     /**
      * Number of threads to use for processing.
diff --git a/mg2d_mpi_test.c b/mg2d_mpi_test.c
index 1bd2784..758940d 100644
--- a/mg2d_mpi_test.c
+++ b/mg2d_mpi_test.c
@@ -238,7 +238,7 @@ int main(int argc, char **argv)
             double rhs = 0.0;
 
             for (int i = 0; i < MG2D_DIFF_COEFF_NB; i++) {
-                ctx->diff_coeffs[i][ctx->diff_coeffs_stride * y + x] = pde_coeffs[i];
+                ctx->diff_coeffs[i]->data[ctx->diff_coeffs[i]->stride * y + x] = pde_coeffs[i];
                 rhs += pde_coeffs[i] * sol[i](x_coord, y_coord);
             }
 
diff --git a/mg2d_test.c b/mg2d_test.c
index fc1d7dd..ec21d7d 100644
--- a/mg2d_test.c
+++ b/mg2d_test.c
@@ -157,7 +157,7 @@ int main(int argc, char **argv)
             double rhs = 0.0;
 
             for (int i = 0; i < MG2D_DIFF_COEFF_NB; i++) {
-                ctx->diff_coeffs[i][ctx->diff_coeffs_stride * y + x] = pde_coeffs[i];
+                ctx->diff_coeffs[i]->data[ctx->diff_coeffs[i]->stride * y + x] = pde_coeffs[i];
                 rhs += pde_coeffs[i] * sol[i](x_coord, y_coord);
             }