mg2d: add API for interpolating an initial guess from a provided grid

5 files changed, 140 insertions, 12 deletions

diff --git a/libmg2d.v b/libmg2d.v
index 3e4fb25..b4901a6 100644
--- a/libmg2d.v
+++ b/libmg2d.v
@@ -1,4 +1,4 @@
-LIBMG2D_11 {
+LIBMG2D_12 {
     global: mg2d_*;
     local: *;
 };
diff --git a/mg2d.c b/mg2d.c
index 03a1759..c9070fb 100644
--- a/mg2d.c
+++ b/mg2d.c
@@ -70,6 +70,8 @@ struct MG2DInternal {
     NDArray *rhs;
     NDArray *diff_coeffs[MG2D_DIFF_COEFF_NB];
 
+    GridTransferContext *transfer_init;
+
     int cpuflags;
 
     int64_t time_solve;
@@ -319,10 +321,25 @@ static int mg_levels_init(MG2DContext *ctx)
     cur  = priv->root;
     prev = NULL;
 
-    mg2di_ndarray_copy(cur->solver->u,   priv->u);
-    mg2di_ndarray_copy(cur->solver->rhs, priv->rhs);
-    for (int i = 0; i < ARRAY_ELEMS(ctx->diff_coeffs); i++)
-        mg2di_ndarray_copy(cur->solver->diff_coeffs[i], priv->diff_coeffs[i]);
+    if (priv->u) {
+        mg2di_ndarray_copy(cur->solver->u,   priv->u);
+        mg2di_ndarray_free(&priv->u);
+        ctx->u = cur->solver->u->data;
+        ctx->u_stride = cur->solver->u->stride[0];
+
+        mg2di_ndarray_copy(cur->solver->rhs, priv->rhs);
+        mg2di_ndarray_free(&priv->rhs);
+        ctx->rhs = cur->solver->rhs->data;
+        ctx->rhs_stride = cur->solver->rhs->stride[0];
+
+        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;
+        }
+        ctx->diff_coeffs_stride = cur->solver->diff_coeffs[0]->stride[0];
+
+    }
 
     while (cur) {
         if (!prev) {
@@ -356,6 +373,7 @@ static int mg_levels_init(MG2DContext *ctx)
 
         cur->solver->logger = priv->logger;
         cur->solver->cpuflags = priv->cpuflags;
+        cur->solver->tp       = priv->tp;
 
         cur->solver->fd_stencil = ctx->fd_stencil;
         if (cur->solver->solver_type == EGS_SOLVER_RELAXATION) {
@@ -449,18 +467,12 @@ static int mg_levels_init(MG2DContext *ctx)
 static int threadpool_init(MG2DContext *ctx)
 {
     MG2DInternal *priv = ctx->priv;
-    MG2DLevel *level = priv->root;
     int ret;
 
     ret = tp_init(&priv->tp, ctx->nb_threads);
     if (ret < 0)
         return ret;
 
-    while (level) {
-        level->solver->tp = priv->tp;
-        level = level->child;
-    }
-
     return 0;
 }
 
@@ -519,7 +531,7 @@ int mg2d_solve(MG2DContext *ctx)
             mg2di_log(&priv->logger, MG2D_LOG_INFO, "converged on iteration %d, residual %g\n",
                       i, res_cur);
 
-            mg2di_ndarray_copy(priv->u, s_root->u);
+            //mg2di_ndarray_copy(priv->u, s_root->u);
 
             priv->time_solve += gettime() - time_start;
             priv->count_solve++;
@@ -741,6 +753,8 @@ void mg2d_solver_free(MG2DContext **pctx)
 
     tp_free(&ctx->priv->tp);
 
+    mg2di_gt_free(&ctx->priv->transfer_init);
+
     mg2di_ndarray_free(&ctx->priv->u);
     mg2di_ndarray_free(&ctx->priv->rhs);
     for (int i = 0; i < ARRAY_ELEMS(ctx->priv->diff_coeffs); i++)
@@ -856,3 +870,60 @@ unsigned int mg2d_max_fd_stencil(void)
 {
     return FD_STENCIL_MAX;
 }
+
+int mg2d_init_guess(MG2DContext *ctx, const double *src,
+                    ptrdiff_t src_stride,
+                    const size_t src_size[2],
+                    const double src_step[2])
+{
+    MG2DInternal *priv = ctx->priv;
+    NDArray *a_src;
+    int ret;
+
+    if (!priv->tp) {
+        ret = threadpool_init(ctx);
+        if (ret < 0)
+            return ret;
+    }
+
+    if (priv->transfer_init &&
+        (priv->transfer_init->src.size[0] != src_size[0] ||
+         priv->transfer_init->src.size[1] != src_size[1] ||
+         fabs(priv->transfer_init->src.step[0] - src_step[0]) > 1e-15 ||
+         fabs(priv->transfer_init->src.step[1] - src_step[1]) > 1e-15)) {
+        mg2di_gt_free(&priv->transfer_init);
+    }
+
+    if (!priv->transfer_init) {
+        priv->transfer_init = mg2di_gt_alloc(GRID_TRANSFER_LAGRANGE_3);
+        if (!priv->transfer_init)
+            return -ENOMEM;
+
+        priv->transfer_init->tp       = priv->tp;
+        priv->transfer_init->cpuflags = priv->cpuflags;
+
+        priv->transfer_init->src.size[0] = src_size[0];
+        priv->transfer_init->src.size[1] = src_size[1];
+        priv->transfer_init->src.step[0] = src_step[0];
+        priv->transfer_init->src.step[1] = src_step[1];
+
+        priv->transfer_init->dst.size[0] = ctx->domain_size;
+        priv->transfer_init->dst.size[1] = ctx->domain_size;
+        priv->transfer_init->dst.step[0] = ctx->step[0];
+        priv->transfer_init->dst.step[1] = ctx->step[1];
+
+        ret = mg2di_gt_init(priv->transfer_init);
+        if (ret < 0)
+            return ret;
+    }
+
+    ret = mg2di_ndarray_wrap(&a_src, 2, src_size, src,
+                             (ptrdiff_t [2]){ src_stride, 1 });
+    if (ret < 0)
+        return ret;
+
+    ret = mg2di_gt_transfer(priv->transfer_init, priv->u ? priv->u : priv->root->solver->u, a_src);
+
+    mg2di_ndarray_free(&a_src);
+    return ret;
+}
diff --git a/mg2d.h b/mg2d.h
index 6452526..fd11906 100644
--- a/mg2d.h
+++ b/mg2d.h
@@ -195,4 +195,9 @@ void mg2d_print_stats(MG2DContext *ctx, const char *prefix);
  */
 unsigned int mg2d_max_fd_stencil(void);
 
+int mg2d_init_guess(MG2DContext *ctx, const double *src,
+                    ptrdiff_t src_stride,
+                    const size_t src_size[2],
+                    const double src_step[2]);
+
 #endif /* MG2D_H */
diff --git a/ndarray.c b/ndarray.c
index a733955..504bfb7 100644
--- a/ndarray.c
+++ b/ndarray.c
@@ -168,6 +168,54 @@ fail:
     return ret;
 }
 
+int mg2di_ndarray_wrap(NDArray **result, unsigned int dims,
+                       const size_t * const size, double *data,
+                       const ptrdiff_t * const stride)
+{
+    NDArray *a;
+    int ret;
+
+    if (!dims || dims >= INT_MAX)
+        return -EINVAL;
+
+    a = calloc(1, sizeof(*a));
+    if (!a)
+        return -ENOMEM;
+
+    a->priv = calloc(1, sizeof(*a->priv));
+    if (!a->priv) {
+        free(a);
+        return -ENOMEM;
+    }
+
+    a->priv->shape = calloc(dims, sizeof(*a->priv->shape));
+    if (!a->priv->shape) {
+        ret = -ENOMEM;
+        goto fail;
+    }
+    memcpy(a->priv->shape, size, dims * sizeof(*a->priv->shape));
+
+    a->priv->stride = calloc(dims, sizeof(*a->priv->stride));
+    if (!a->priv->stride) {
+        ret = -ENOMEM;
+        goto fail;
+    }
+    memcpy(a->priv->stride, stride, dims * sizeof(*a->priv->shape));
+
+    a->dims   = dims;
+
+    a->data   = data;
+    a->shape  = a->priv->shape;
+    a->stride = a->priv->stride;
+
+    *result = a;
+
+    return 0;
+fail:
+    mg2di_ndarray_free(&a);
+    return ret;
+}
+
 void mg2di_ndarray_free(NDArray **pa)
 {
     NDArray *a = *pa;
diff --git a/ndarray.h b/ndarray.h
index c0d964b..9e1b8ba 100644
--- a/ndarray.h
+++ b/ndarray.h
@@ -57,6 +57,10 @@ int mg2di_ndarray_alloc(NDArray **result, unsigned int dims,
                         const size_t * const size, unsigned int flags);
 void mg2di_ndarray_free(NDArray **a);
 
+int mg2di_ndarray_wrap(NDArray **result, unsigned int dims,
+                       const size_t * const size, double *data,
+                       const ptrdiff_t * const strides);
+
 int mg2di_ndarray_slice(NDArray **result, NDArray *src,
                         const Slice * const slice);