Implement the multigrid scheme.

1 files changed, 421 insertions, 0 deletions

diff --git a/mg2d.c b/mg2d.c
new file mode 100644
index 0000000..1e605b8
--- /dev/null
+++ b/mg2d.c
@@ -0,0 +1,421 @@
+/*
+ * Multigrid solver for a 2nd order 2D linear PDE.
+ * Copyright 2018 Anton Khirnov <anton@khirnov.net>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include <errno.h>
+#include <math.h>
+#include <stdlib.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <string.h>
+
+#include "common.h"
+#include "ell_relax.h"
+#include "log.h"
+#include "mg2d.h"
+
+#define LEVELS_MAX 16
+
+typedef struct MG2DLevel {
+    unsigned int  depth;
+
+    EllRelaxContext *solver;
+
+    double   *prolong_tmp;
+    ptrdiff_t prolong_tmp_stride;
+
+    struct MG2DLevel *child;
+} MG2DLevel;
+
+struct MG2DInternal {
+    MG2DLogger logger;
+
+    MG2DLevel *root;
+};
+
+static double findmax(double *arr, size_t len)
+{
+    double ret = 0.0;
+    for (size_t i = 0; i < len; i++) {
+        double val = fabs(*arr++);
+        if (val > ret)
+            ret = val;
+    }
+    return ret;
+}
+
+static void log_callback(MG2DLogger *log, int level, const char *fmt, va_list vl)
+{
+    MG2DContext *ctx = log->opaque;
+    ctx->log_callback(ctx, level, fmt, vl);
+}
+
+static void mg_restrict(EllRelaxContext *coarse, double *dst, ptrdiff_t dst_stride,
+                        EllRelaxContext *fine, const double *src, ptrdiff_t src_stride)
+{
+    for (size_t j = 0; j < coarse->domain_size[1]; j++)
+        for (size_t i = 0; i < coarse->domain_size[0]; i++) {
+            const ptrdiff_t idx_dst = j * dst_stride + i;
+
+            const size_t fine_i = i * 2;
+            const size_t fine_j = j * 2;
+            const ptrdiff_t idx_src = fine_j * src_stride + fine_i;
+
+            //dst[idx_dst] = src[idx_src];
+            dst[idx_dst] = 0.25 * src[idx_src] + 0.125 * (src[idx_src + 1] + src[idx_src - 1] + src[idx_src + src_stride] + src[idx_src - src_stride])
+                           + 0.0625 * (src[idx_src + 1 + src_stride] + src[idx_src - 1 + src_stride] + src[idx_src + 1 - src_stride] + src[idx_src - 1 - src_stride]);
+        }
+}
+
+static void mg_prolongate(EllRelaxContext *fine, double *dst, ptrdiff_t dst_stride,
+                          EllRelaxContext *coarse, const double *src, ptrdiff_t src_stride)
+{
+    for (size_t j = 0; j < coarse->domain_size[1]; j++)
+        for (size_t i = 0; i < coarse->domain_size[0]; i++) {
+            const ptrdiff_t idx_src = j * src_stride + i;
+            const ptrdiff_t idx_dst = 2 * (j * dst_stride + i);
+            dst[idx_dst] = src[idx_src];
+        }
+
+    for (size_t j = 0; j < fine->domain_size[1]; j += 2)
+        for (size_t i = 1; i < fine->domain_size[0]; i += 2) {
+            const ptrdiff_t idx_dst = j * dst_stride + i;
+            dst[idx_dst] = 0.5 * (dst[idx_dst - 1] + dst[idx_dst + 1]);
+        }
+
+    for (size_t j = 1; j < fine->domain_size[1]; j += 2)
+        for (size_t i = 0; i < fine->domain_size[0]; i++) {
+            const ptrdiff_t idx = j * dst_stride + i;
+            dst[idx] = 0.5 * (dst[idx - dst_stride] + dst[idx + dst_stride]);
+        }
+}
+
+static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level)
+{
+    int ret;
+
+    /* on the refined levels, use zero as the initial guess for the
+     * solution (correction for the upper level) */
+    if (level->depth > 0) {
+        memset(level->solver->u, 0, sizeof(*level->solver->u) * level->solver->u_stride *
+                level->solver->domain_size[1]);
+
+        /* re-init the current-level solver (re-calc the residual) */
+        ret = mg2di_ell_relax_init(level->solver);
+        if (ret < 0)
+            return ret;
+    }
+
+    /* handle coarsest grid */
+    if (!level->child) {
+        for (int j = 0; j < 16; j++) {
+            ret = mg2di_ell_relax_step(level->solver);
+            if (ret < 0)
+                return ret;
+        }
+        goto finish;
+    }
+
+    for (int i = 0; i < ctx->nb_cycles; i++) {
+        /* pre-restrict relaxation */
+        for (int j = 0; j < ctx->nb_relax_pre; j++) {
+            ret = mg2di_ell_relax_step(level->solver);
+            if (ret < 0)
+                return ret;
+        }
+
+        /* restrict the residual as to the coarser-level rhs */
+        mg_restrict(level->child->solver, level->child->solver->rhs, level->child->solver->rhs_stride,
+                    level->solver, level->solver->residual, level->solver->residual_stride);
+
+
+        /* solve on the coarser level */
+        ret = mg_solve_subgrid(ctx, level->child);
+        if (ret < 0)
+            return ret;
+
+        /* prolongate the coarser-level correction */
+        mg_prolongate(level->solver, level->prolong_tmp, level->prolong_tmp_stride,
+                      level->child->solver, level->child->solver->u, level->child->solver->u_stride);
+
+        /* apply the correction */
+        for (size_t idx1 = 0; idx1 < level->solver->domain_size[1]; idx1++)
+            for (size_t idx0 = 0; idx0 < level->solver->domain_size[0]; idx0++) {
+                const ptrdiff_t idx_dst = idx1 * level->solver->u_stride + idx0;
+                const ptrdiff_t idx_src = idx1 * level->prolong_tmp_stride + idx0;
+                level->solver->u[idx_dst] -= level->prolong_tmp[idx_src];
+            }
+
+        /* re-init the current-level solver (re-calc the residual) */
+        ret = mg2di_ell_relax_init(level->solver);
+        if (ret < 0)
+            return ret;
+
+        /* post-correct relaxation */
+        for (int j = 0; j < ctx->nb_relax_pre; j++) {
+            ret = mg2di_ell_relax_step(level->solver);
+            if (ret < 0)
+                return ret;
+        }
+    }
+
+finish:
+    return 0;
+}
+
+static int mg_levels_init(MG2DContext *ctx)
+{
+    MG2DInternal *priv = ctx->priv;
+    MG2DLevel *cur, *prev;
+
+    cur  = priv->root;
+    prev = NULL;
+    while (cur) {
+        /* Set the equation coefficients. */
+        if (prev) {
+            for (int i = 0; i < ARRAY_ELEMS(prev->solver->diff_coeffs); i++) {
+                mg_restrict(cur->solver, cur->solver->diff_coeffs[i], cur->solver->diff_coeffs_stride,
+                            prev->solver, prev->solver->diff_coeffs[i], prev->solver->diff_coeffs_stride);
+            }
+        }
+
+        for (int i = 0; i < ARRAY_ELEMS(cur->solver->boundaries); i++) {
+            MG2DBoundary *bsrc = ctx->boundaries[i];
+            EllRelaxBoundary *bdst = &cur->solver->boundaries[i];
+            switch (bsrc->type) {
+            case MG2D_BC_TYPE_FIXVAL:
+                bdst->type = ELL_RELAX_BC_TYPE_FIXVAL;
+                if (!prev)
+                    memcpy(bdst->val, bsrc->val, sizeof(*bsrc->val) * bsrc->val_len);
+                else
+                    memset(bdst->val, 0, sizeof(*bdst->val) * bdst->val_len);
+                break;
+            case MG2D_BC_TYPE_FIXDIFF:
+                bdst->type = ELL_RELAX_BC_TYPE_FIXDIFF;
+                memset(bdst->val, 0, sizeof(*bdst->val) * bdst->val_len);
+                break;
+            default:
+                return -ENOSYS;
+            }
+        }
+
+        cur->solver->logger = priv->logger;
+
+        cur->solver->step[0] = prev ? 2.0 * prev->solver->step[0] : ctx->step[0];
+        cur->solver->step[1] = prev ? 2.0 * prev->solver->step[1] : ctx->step[1];
+
+        cur->solver->fd_stencil = ctx->fd_stencil;
+
+        prev = cur;
+        cur  = cur->child;
+    }
+
+    return 0;
+}
+
+int mg2d_solve(MG2DContext *ctx)
+{
+    MG2DInternal *priv = ctx->priv;
+    double res_prev, res_cur;
+    int ret;
+
+    ret = mg_levels_init(ctx);
+    if (ret < 0)
+        return ret;
+
+    ret = mg2di_ell_relax_init(priv->root->solver);
+    if (ret < 0)
+        return ret;
+
+    res_prev = findmax(priv->root->solver->residual,
+                       priv->root->solver->residual_stride * ctx->domain_size);
+
+    for (int i = 0; i < ctx->maxiter; i++) {
+        mg_solve_subgrid(ctx, priv->root);
+
+        res_cur = findmax(priv->root->solver->residual,
+                          priv->root->solver->residual_stride * ctx->domain_size);
+
+        if (res_cur < ctx->tol) {
+            mg2di_log(&priv->logger, 0, "converged on iteration %d, residual %g\n",
+                      i, res_cur);
+            return 0;
+        }
+
+        mg2di_log(&priv->logger, 0, "finished toplevel iteration %d, residual %g -> %g (%g)\n",
+                  i, res_prev, res_cur, res_prev / res_cur);
+        res_prev = res_cur;
+    }
+
+    mg2di_log(&priv->logger, 0, "The solver failed to converge\n");
+    return -EDOM;
+}
+
+static void mg_level_free(MG2DLevel **plevel)
+{
+    MG2DLevel *level = *plevel;
+
+    if (!level)
+        return;
+
+    free(level->prolong_tmp);
+    mg2di_ell_relax_free(&level->solver);
+
+    free(level);
+    *plevel = NULL;
+}
+
+static MG2DLevel *mg_level_alloc(const size_t domain_size)
+{
+    MG2DLevel *level;
+    int ret;
+
+    level = calloc(1, sizeof(*level));
+    if (!level)
+        return NULL;
+
+    ret = posix_memalign((void**)&level->prolong_tmp, 32, sizeof(*level->prolong_tmp) * SQR(domain_size));
+    if (ret != 0)
+        goto fail;
+    level->prolong_tmp_stride = domain_size;
+
+    level->solver = mg2di_ell_relax_alloc((size_t [2]){domain_size, domain_size});
+    if (!level->solver)
+        goto fail;
+
+    return level;
+fail:
+    mg_level_free(&level);
+    return NULL;
+}
+
+static int mg_levels_alloc(MG2DContext *ctx, size_t domain_size)
+{
+    MG2DInternal *priv = ctx->priv;
+    MG2DLevel *cur;
+
+    priv->root = mg_level_alloc(domain_size);
+    if (!priv->root)
+        return -ENOMEM;
+
+    cur = priv->root;
+    domain_size = (domain_size >> 1) + 1;
+    for (int i = 0; i < LEVELS_MAX && domain_size >= 4; i++) {
+        cur->child = mg_level_alloc(domain_size);
+        if (!cur->child)
+            return -ENOMEM;
+        cur->child->depth = i + 1;
+
+        cur = cur->child;
+        domain_size = (domain_size >> 1) + 1;
+    }
+
+    return 0;
+}
+
+static void log_default_callback(const MG2DContext *ctx, int level, const char *fmt, va_list vl)
+{
+    vfprintf(stderr, fmt, vl);
+}
+
+MG2DContext *mg2d_solver_alloc(size_t log2_gridsize)
+{
+    MG2DContext *ctx;
+    MG2DInternal *priv;
+    size_t domain_size;
+    int ret;
+
+    ctx = calloc(1, sizeof(*ctx));
+    if (!ctx)
+        return NULL;
+
+    ctx->priv = calloc(1, sizeof(*ctx->priv));
+    if (!ctx->priv)
+        goto fail;
+    priv = ctx->priv;
+
+    priv->logger.log    = log_callback;
+    priv->logger.opaque = ctx;
+
+    for (int i = 0; i < ARRAY_ELEMS(ctx->boundaries); i++) {
+        ctx->boundaries[i] = calloc(1, sizeof(*ctx->boundaries[i]));
+        if (!ctx->boundaries[i])
+            goto fail;
+
+        ret = posix_memalign((void**)&ctx->boundaries[i]->val, 32,
+                             sizeof(*ctx->boundaries[i]->val) * ctx->domain_size);
+        if (ret != 0)
+            goto fail;
+        ctx->boundaries[i]->val_len = ctx->domain_size;
+    }
+
+    if (!((SIZE_MAX - 1) >> log2_gridsize))
+        goto fail;
+    domain_size = (1 << log2_gridsize) + 1;
+
+    ret = mg_levels_alloc(ctx, domain_size);
+    if (ret < 0)
+        goto fail;
+
+    ctx->domain_size = domain_size;
+
+    ctx->maxiter       = 16;
+    ctx->tol           = 1e-12;
+    ctx->nb_cycles     = 1;
+    ctx->nb_relax_pre  = 1;
+    ctx->nb_relax_post = 1;
+    ctx->log_callback  = log_default_callback;
+
+    ctx->u          = priv->root->solver->u;
+    ctx->u_stride   = priv->root->solver->u_stride;
+    ctx->rhs        = priv->root->solver->rhs;
+    ctx->rhs_stride = priv->root->solver->rhs_stride;
+
+    for (int i = 0; i < ARRAY_ELEMS(ctx->diff_coeffs); i++)
+        ctx->diff_coeffs[i] = priv->root->solver->diff_coeffs[i];
+    ctx->diff_coeffs_stride = priv->root->solver->diff_coeffs_stride;
+
+    return ctx;
+fail:
+    mg2d_solver_free(&ctx);
+    return NULL;
+}
+
+void mg2d_solver_free(MG2DContext **pctx)
+{
+    MG2DContext *ctx = *pctx;
+
+    if (!ctx)
+        return;
+
+    while (ctx->priv->root) {
+        MG2DLevel *next = ctx->priv->root->child;
+        mg_level_free(&ctx->priv->root);
+        ctx->priv->root = next;
+    }
+
+    free(ctx->priv);
+
+    for (int i = 0; i < ARRAY_ELEMS(ctx->boundaries); i++) {
+        if (ctx->boundaries[i])
+            free(ctx->boundaries[i]->val);
+        free(ctx->boundaries[i]);
+    }
+
+    free(ctx);
+    *pctx = NULL;
+}