mg2d: add support for adaptive step sizeadaptive_step

3 files changed, 139 insertions, 27 deletions

diff --git a/mg2d.c b/mg2d.c
index 103f514..4b9af11 100644
--- a/mg2d.c
+++ b/mg2d.c
@@ -37,6 +37,7 @@
 #include "mg2d_boundary.h"
 #include "mg2d_boundary_internal.h"
 #include "mg2d_constants.h"
+#include "step_control.h"
 #include "timer.h"
 #include "transfer.h"
 
@@ -83,6 +84,8 @@ typedef struct MG2DLevel {
     char          log_prefix[32];
     MG2DInternal *priv;
 
+    StepControl *sc;
+
     /* timings */
     int64_t count_cycles;
     Timer   timer_solve;
@@ -112,6 +115,7 @@ struct MG2DInternal {
     MG2DLevel *root;
 
     NDArray *u;
+    NDArray *u_exterior;
     NDArray *rhs;
     NDArray *diff_coeffs_base[MG2D_DIFF_COEFF_NB];
     NDArray *diff_coeffs_tmp[MG2D_DIFF_COEFF_NB];
@@ -233,6 +237,7 @@ static int mg_prolong_u(MG2DContext *ctx, MG2DLevel *l)
     return 0;
 }
 
+#define SUBGRID_DIVERGE 1
 static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level, int allow_exact)
 {
     enum EGSType solve_type = (allow_exact && level->dg->nb_components == 1 &&
@@ -269,7 +274,7 @@ static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level, int allow_exact)
         goto finish;
     }
 
-    for (int i = 0; i < ctx->nb_cycles; i++) {
+    for (int cycle = 0; cycle < ctx->nb_cycles; cycle++) {
         double res_prev;
 
         /* pre-restrict relaxation */
@@ -282,17 +287,19 @@ static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level, int allow_exact)
         }
 
         if (level->depth < ctx->priv->dh.nb_levels - 1) {
-            /* restrict the residual as to the coarser-level rhs */
-            ret = mg_restrict_rhs(ctx, level);
-            if (ret < 0)
-                return ret;
-
-            /* solve on the coarser level */
-            if (level->child) {
-                ret = mg_solve_subgrid(ctx, level->child, 1);
+            do {
+                /* restrict the residual as to the coarser-level rhs */
+                ret = mg_restrict_rhs(ctx, level);
                 if (ret < 0)
                     return ret;
-            }
+
+                /* solve on the coarser level */
+                if (level->child) {
+                    ret = mg_solve_subgrid(ctx, level->child, 1);
+                    if (ret < 0)
+                        return ret;
+                }
+            } while (ret == SUBGRID_DIVERGE);
 
             /* prolongate the coarser-level correction */
             ret = mg_prolong_u(ctx, level);
@@ -318,10 +325,50 @@ static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level, int allow_exact)
         }
 
         /* post-correct relaxation */
-        for (int j = 0; j < ctx->nb_relax_post; j++) {
-            ret = mg_solve(ctx, level, solve_type, "post-relax", 0);
-            if (ret < 0)
-                return ret;
+        {
+            double conv_history[4][2] = { { 0. } };
+            int conv_hist_next = 0;
+
+            double min_fact = DBL_MAX;
+            int min_fact_idx = -1;
+
+            for (int j = 0; j < ctx->nb_relax_post; j++) {
+                conv_history[conv_hist_next][0] = level->solver->residual_max;
+
+                ret = mg_solve(ctx, level, solve_type, "post-relax", 0);
+                if (ret < 0)
+                    return ret;
+
+                conv_history[conv_hist_next][1] = level->solver->residual_max;
+                conv_hist_next = (conv_hist_next + 1) % ARRAY_ELEMS(conv_history);
+            }
+
+            // find the minimum convergence factor in the history window
+            for (int i = 0; i < ARRAY_ELEMS(conv_history); i++) {
+                double *h = conv_history[i], f;
+                if (h[0] == 0.0)
+                    continue;
+                f = h[0] / h[1];
+                if (f < min_fact) {
+                    min_fact = f;
+                    min_fact_idx = i;
+                }
+            }
+
+            // use the minimum convergence factor as input for step control
+            if (ctx->adaptive_step &&
+                level->solver->residual_max > ctx->tol && min_fact_idx >= 0) {
+                double *h = conv_history[min_fact_idx];
+
+                ret = sc_step_confirm(level->sc, level->solver->relax->relax_factor,
+                                     h[0], h[1]);
+                if (ret < 0)
+                    return ret;
+
+                level->solver->relax->relax_factor = sc_step_get(level->sc);
+                if (ret)
+                    return SUBGRID_DIVERGE;
+            }
         }
 
         level->count_cycles++;
@@ -791,9 +838,6 @@ static int mg_levels_init(MG2DContext *ctx)
         ret = ndarray_copy(cur->solver->u,   priv->u);
         if (ret < 0)
             return ret;
-        ndarray_free(&priv->u);
-        ctx->u = cur->solver->u->data;
-        ctx->u_stride = cur->solver->u->stride[0];
     }
 
     if (priv->rhs) {
@@ -906,7 +950,12 @@ static int mg_levels_init(MG2DContext *ctx)
 
     cur = priv->root;
     while (cur) {
-        cur->solver->relax->relax_factor = ctx->cfl_factor;
+        if (ctx->adaptive_step) {
+            sc_init(cur->sc, ctx->cfl_factor, ctx->cfl_factor * 1e-5);
+            cur->solver->relax->relax_factor = sc_step_get(cur->sc);
+        } else {
+            cur->solver->relax->relax_factor = ctx->cfl_factor;
+        }
         cur->solver->relax->relax_multiplier = 1.0 / (diff2_max + cur->solver->step[0] * cur->solver->step[1] *
                                                  diff0_max / 8.0);
 
@@ -1135,9 +1184,15 @@ int mg2d_solve(MG2DContext *ctx)
     res_prev = res_orig;
 
     for (int i = 0; i < ctx->maxiter; i++) {
-        ret = mg_solve_subgrid(ctx, root, 1);
-        if (ret < 0)
-            goto fail;
+        do {
+            ret = mg_solve_subgrid(ctx, root, 1);
+            if (ret < 0)
+                goto fail;
+            if (ret == SUBGRID_DIVERGE) {
+                ndarray_copy(s_root->u,   priv->u);
+                mg2di_egs_init(s_root, 0);
+            }
+        } while (ret == SUBGRID_DIVERGE);
 
         res_cur = s_root->residual_max;
 
@@ -1145,14 +1200,13 @@ int mg2d_solve(MG2DContext *ctx)
             mg2di_log(&priv->logger, MG2D_LOG_INFO, "converged on iteration %d, residual %g\n",
                       i, res_cur);
 
-            //ndarray_copy(priv->u, s_root->u);
-
             goto finish;
         }
+        ndarray_copy(priv->u_exterior, s_root->u_exterior);
 
         mg2di_log(&priv->logger, MG2D_LOG_VERBOSE,
-                  "finished toplevel iteration %d, residual %g -> %g (%g)\n",
-                  i, res_prev, res_cur, res_prev / res_cur);
+                  "finished toplevel iteration %d, residual %g -> %g (%g) %g\n",
+                  i, res_prev, res_cur, res_prev / res_cur, s_root->relax->relax_factor);
         if (res_cur / res_prev > 1e1 || res_cur / res_orig > 1e3) {
             mg2di_log(&priv->logger, MG2D_LOG_ERROR, "A multigrid iteration diverged\n");
             ret = MG2D_ERR_DIVERGE;
@@ -1167,6 +1221,7 @@ int mg2d_solve(MG2DContext *ctx)
 fail:
     mg2di_log(&priv->logger, MG2D_LOG_ERROR, "The solver failed to converge\n");
 finish:
+    ndarray_copy(priv->u_exterior, s_root->u_exterior);
     ctx->residual_max = s_root->residual_max;
     mg2di_timer_stop(&priv->timer_solve);
     return ret;
@@ -1227,6 +1282,8 @@ static void mg_level_free(MG2DLevel **plevel)
     ndarray_free(&level->prolong_tmp_base);
     mg2di_egs_free(&level->solver);
 
+    sc_free(&level->sc);
+
     mg2di_gt_free(&level->transfer_restrict);
     mg2di_gt_free(&level->transfer_prolong);
 
@@ -1550,6 +1607,12 @@ static int mg_levels_alloc(MG2DContext *ctx)
         snprintf(level->log_prefix + level->depth, MAX(0, sizeof(level->log_prefix) - level->depth),
                  "[%d]", depth);
 
+        ret = sc_alloc(&level->sc, level->logger);
+        if (ret < 0) {
+            mg_level_free(&level);
+            return ret;
+        }
+
         *dst = level;
         dst  = &level->child;
         comm_parent = comm_cur;
@@ -1568,6 +1631,45 @@ static int mg_levels_alloc(MG2DContext *ctx)
         }
     }
 
+    // allocate the user-facing array for the solution of the same size as the topmost level's
+    {
+        const NDArray *src_exterior = priv->root->solver->u_exterior;
+        const NDArray *src_interior = priv->root->solver->u;
+        ptrdiff_t offset = src_interior->data - src_exterior->data;
+
+        NDArray *tmp;
+        NDASlice slice[2];
+
+        ret = ndarray_alloc(&tmp, 2, src_exterior->shape, 0);
+        if (ret < 0)
+            return ret;
+
+        ndarray_copy(tmp, src_exterior);
+
+        ndarray_free(&priv->u);
+        ctx->u = NULL;
+
+        ndarray_free(&priv->u_exterior);
+
+        priv->u_exterior = tmp;
+        tmp              = NULL;
+
+        for (int i = 0; i < ARRAY_ELEMS(slice); i++) {
+            size_t ext_size = src_exterior->shape[i];
+            size_t int_size = src_interior->shape[i];
+
+            slice[i].start = (offset / src_exterior->stride[i]) % ext_size;
+            slice[i].end   = -(ext_size - int_size - slice[i].start);
+            slice[i].step  = 1;
+        }
+
+        ret = ndarray_slice(&priv->u, priv->u_exterior, slice);
+        if (ret < 0)
+            return ret;
+        ctx->u        = priv->u->data;
+        ctx->u_stride = priv->u->stride[0];
+    }
+
     return ret;
 }
 
@@ -1614,10 +1716,13 @@ static MG2DContext *solver_alloc(DomainGeometry *dg, unsigned int local_componen
         }
     }
 
-    ret = ndarray_alloc(&priv->u, 2, (size_t [2]){ domain_size[1], domain_size[0] },
+    ret = ndarray_alloc(&priv->u_exterior, 2, (size_t [2]){ domain_size[1], domain_size[0] },
                               NDARRAY_ALLOC_ZERO);
     if (ret < 0)
         goto fail;
+    ret = ndarray_slice(&priv->u, priv->u_exterior, (const NDASlice [2]){ NDASLICE_NULL, NDASLICE_NULL });
+    if (ret < 0)
+        goto fail;
     ctx->u        = priv->u->data;
     ctx->u_stride = priv->u->stride[0];
 
@@ -1835,6 +1940,7 @@ void mg2d_solver_free(MG2DContext **pctx)
     mg2di_gt_free(&ctx->priv->transfer_init);
 
     ndarray_free(&ctx->priv->u);
+    ndarray_free(&ctx->priv->u_exterior);
     ndarray_free(&ctx->priv->rhs);
     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++)
@@ -2038,7 +2144,7 @@ int mg2d_init_guess(MG2DContext *ctx, const double *src,
     if (ret < 0)
         return ret;
 
-    ret = mg2di_gt_transfer(priv->transfer_init, priv->u ? priv->u : priv->root->solver->u, a_src);
+    ret = mg2di_gt_transfer(priv->transfer_init, priv->u, a_src);
 
     ndarray_free(&a_src);
     return ret;
diff --git a/mg2d.h b/mg2d.h
index 8f48361..56e129f 100644
--- a/mg2d.h
+++ b/mg2d.h
@@ -246,6 +246,11 @@ typedef struct MG2DContext {
      * Set by mg2d_solve() if it returns 0 or MG2D_ERR_MAXITER_REACHED.
      */
     double residual_max;
+
+    /**
+     * Use adaptive stepping for relaxation.
+     */
+    int adaptive_step;
 } MG2DContext;
 
 /**
diff --git a/mg2d_test.c b/mg2d_test.c
index ec21d7d..f840363 100644
--- a/mg2d_test.c
+++ b/mg2d_test.c
@@ -120,6 +120,7 @@ int main(int argc, char **argv)
     ctx->nb_relax_post = 2;
     ctx->tol = TOL / (ctx->step[0] * ctx->step[1]);
     ctx->nb_threads = 1;
+    ctx->adaptive_step = 1;
     ctx->log_level = MG2D_LOG_INFO;
 
     for (int bnd_loc = 0; bnd_loc < ARRAY_ELEMS(ctx->boundaries); bnd_loc++) {