mg2d_interp tmp

2 files changed, 137 insertions, 0 deletions

diff --git a/mg2d.c b/mg2d.c
index ca821d2..3901b9f 100644
--- a/mg2d.c
+++ b/mg2d.c
@@ -119,6 +119,7 @@ struct MG2DInternal {
     NDArray *dc_bnd_val[MG2D_DIFF_COEFF_NB][4];
 
     GridTransferContext *transfer_init;
+    GridTransferContext *transfer_interp;
 
     int cpuflags;
 
@@ -1804,6 +1805,7 @@ void mg2d_solver_free(MG2DContext **pctx)
     tp_free(&ctx->priv->tp);
 
     mg2di_gt_free(&ctx->priv->transfer_init);
+    mg2di_gt_free(&ctx->priv->transfer_interp);
 
     mg2di_ndarray_free(&ctx->priv->u);
     mg2di_ndarray_free(&ctx->priv->rhs);
@@ -2014,3 +2016,120 @@ int mg2d_init_guess(MG2DContext *ctx, const double *src,
     mg2di_ndarray_free(&a_src);
     return ret;
 }
+
+int mg2d_interp_solution(MG2DContext *ctx, double *dst,
+                         ptrdiff_t dst_stride,
+                         const ptrdiff_t dst_start[2],
+                         const size_t    dst_size[2],
+                         const double    dst_step[2])
+{
+    MG2DInternal *priv = ctx->priv;
+    NDArray *a_dst;
+    int ret;
+
+    if (priv->transfer_interp &&
+        (priv->transfer_interp->dst.size[0] != dst_size[0] ||
+         priv->transfer_interp->dst.size[1] != dst_size[1] ||
+         fabs(priv->transfer_interp->dst.step[0] - dst_step[0]) > 1e-15 ||
+         fabs(priv->transfer_interp->dst.step[1] - dst_step[1]) > 1e-15)) {
+        mg2di_gt_free(&priv->transfer_interp);
+    }
+
+    if (!priv->transfer_interp) {
+        priv->transfer_interp = mg2di_gt_alloc(2, GRID_TRANSFER_LAGRANGE_5);
+        if (!priv->transfer_interp)
+            return -ENOMEM;
+
+        priv->transfer_interp->tp       = priv->tp;
+        priv->transfer_interp->cpuflags = priv->cpuflags;
+
+        priv->transfer_interp->dst.start[0] = dst_start[1];
+        priv->transfer_interp->dst.start[1] = dst_start[0];
+        priv->transfer_interp->dst.size[0]  = dst_size[1];
+        priv->transfer_interp->dst.size[1]  = dst_size[0];
+        priv->transfer_interp->dst.step[0]  = dst_step[1];
+        priv->transfer_interp->dst.step[1]  = dst_step[0];
+
+        priv->transfer_interp->src.start[0] = ctx->local_start[1];
+        priv->transfer_interp->src.start[1] = ctx->local_start[0];
+        priv->transfer_interp->src.size[0] = ctx->local_size[1];
+        priv->transfer_interp->src.size[1] = ctx->local_size[0];
+        priv->transfer_interp->src.step[0] = ctx->step[1];
+        priv->transfer_interp->src.step[1] = ctx->step[0];
+
+        ret = mg2di_gt_init(priv->transfer_interp);
+        if (ret < 0)
+            return ret;
+    }
+
+    ret = mg2di_ndarray_wrap(&a_dst, 2, (size_t [2]){ dst_size[1], dst_size[0] }, dst,
+                             (ptrdiff_t [2]){ dst_stride, 1 });
+    if (ret < 0)
+        return ret;
+
+    ret = mg2di_gt_transfer(priv->transfer_interp, a_dst, priv->u ? priv->u : priv->root->solver->u);
+
+    mg2di_ndarray_free(&a_dst);
+    return ret;
+}
+
+int mg2d_interp(MG2DContext *ctx,
+                double *dst,
+                ptrdiff_t dst_stride,
+                const ptrdiff_t dst_start[2],
+                const size_t    dst_size[2],
+                const double    dst_step[2],
+                double *src,
+                ptrdiff_t src_stride,
+                const ptrdiff_t src_start[2],
+                const size_t    src_size[2],
+                const double    src_step[2])
+{
+    MG2DInternal *priv = ctx->priv;
+    GridTransferContext *gt;
+    NDArray *a_dst;
+    NDArray *a_src;
+    int ret;
+
+    gt = mg2di_gt_alloc(2, GRID_TRANSFER_LAGRANGE_5);
+    if (!gt)
+        return -ENOMEM;
+
+    gt->tp       = priv->tp;
+    gt->cpuflags = priv->cpuflags;
+
+    gt->dst.start[0] = dst_start[1];
+    gt->dst.start[1] = dst_start[0];
+    gt->dst.size[0]  = dst_size[1];
+    gt->dst.size[1]  = dst_size[0];
+    gt->dst.step[0]  = dst_step[1];
+    gt->dst.step[1]  = dst_step[0];
+
+    gt->src.start[0] = src_start[1];
+    gt->src.start[1] = src_start[0];
+    gt->src.size[0] = src_size[1];
+    gt->src.size[1] = src_size[0];
+    gt->src.step[0] = src_step[1];
+    gt->src.step[1] = src_step[0];
+
+    ret = mg2di_gt_init(gt);
+    if (ret < 0)
+        return ret;
+
+    ret = mg2di_ndarray_wrap(&a_dst, 2, (size_t [2]){ dst_size[1], dst_size[0] }, dst,
+                             (ptrdiff_t [2]){ dst_stride, 1 });
+    if (ret < 0)
+        return ret;
+
+    ret = mg2di_ndarray_wrap(&a_src, 2, (size_t [2]){ src_size[1], src_size[0] }, src,
+                             (ptrdiff_t [2]){ src_stride, 1 });
+    if (ret < 0)
+        return ret;
+
+    ret = mg2di_gt_transfer(gt, a_dst, a_src);
+
+    mg2di_ndarray_free(&a_dst);
+    mg2di_ndarray_free(&a_src);
+    mg2di_gt_free(&gt);
+    return ret;
+}
diff --git a/mg2d.h b/mg2d.h
index 867d77e..f7897f6 100644
--- a/mg2d.h
+++ b/mg2d.h
@@ -292,4 +292,22 @@ int mg2d_init_guess(MG2DContext *ctx, const double *src,
                     const size_t src_size[2],
                     const double src_step[2]);
 
+int mg2d_interp_solution(MG2DContext *ctx, double *dst,
+                         ptrdiff_t dst_stride,
+                         const ptrdiff_t dst_start[2],
+                         const size_t    dst_size[2],
+                         const double    dst_step[2]);
+
+int mg2d_interp(MG2DContext *ctx,
+                double *dst,
+                ptrdiff_t dst_stride,
+                const ptrdiff_t dst_start[2],
+                const size_t    dst_size[2],
+                const double    dst_step[2],
+                double *src,
+                ptrdiff_t src_stride,
+                const ptrdiff_t src_start[2],
+                const size_t    src_size[2],
+                const double    src_step[2]);
+
 #endif /* MG2D_H */