Add the generic non-linear solver layer.

1 files changed, 300 insertions, 0 deletions

diff --git a/tests/nlsolve.c b/tests/nlsolve.c
new file mode 100644
index 0000000..6d46c9e
--- /dev/null
+++ b/tests/nlsolve.c
@@ -0,0 +1,300 @@
+/*
+ * Copyright 2017 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 <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include "basis.h"
+#include "common.h"
+#include "log.h"
+#include "nlsolve.h"
+#include "pssolve.h"
+
+#define N_X 8
+#define N_Z 4
+#define FUNC_X 4
+#define FUNC_Z 2
+#define TOL 5e-15
+#define ITER 10
+static int
+scalar_linear_func(void *opaque, unsigned int eq_idx,
+                   const unsigned int *colloc_grid_order,
+                   const double * const *colloc_grid,
+                   const double * const (*vars)[PSSOLVE_DIFF_ORDER_NB],
+                   double *dst)
+{
+    BasisSetContext *basis = opaque;
+
+    for (int j = 0; j < colloc_grid_order[1]; j++) {
+        const double z = colloc_grid[1][j];
+
+        for (int i = 0; i < colloc_grid_order[0]; i++) {
+            const double x = colloc_grid[0][i];
+
+            const int idx = j * colloc_grid_order[0] + i;
+
+            const double d2x = vars[0][PSSOLVE_DIFF_ORDER_20][idx];
+            const double d2z = vars[0][PSSOLVE_DIFF_ORDER_02][idx];
+            const double val = vars[0][PSSOLVE_DIFF_ORDER_00][idx];
+
+            const double d2x_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_DIFF2, x, FUNC_X) *
+                                  tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, z, FUNC_Z);
+            const double d2z_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, x, FUNC_X) *
+                                  tdi_basis_eval(basis, BS_EVAL_TYPE_DIFF2, z, FUNC_Z);
+            const double val_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, x, FUNC_X) *
+                                  tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, z, FUNC_Z);
+
+            dst[idx] = (d2x + d2z + val) - (d2x_ex + d2z_ex + val_ex);
+        }
+    }
+
+    return 0;
+}
+
+static int scalar_linear(void)
+{
+    BasisSetContext *basis = NULL;
+    NLSolveContext    *ctx = NULL;
+
+    double coeffs[N_X * N_Z];
+    int ret;
+
+    ret = tdi_basis_init(&basis, BASIS_FAMILY_SB_EVEN, 1.0);
+    if (ret < 0)
+        goto finish;
+
+    ret = tdi_nlsolve_context_alloc(&ctx, 1);
+    if (ret < 0)
+        return ret;
+
+    ctx->basis[0][0] = basis;
+    ctx->basis[0][1] = basis;
+
+    ctx->solve_order[0][0] = N_X;
+    ctx->solve_order[0][1] = N_Z;
+
+    ctx->maxiter = ITER;
+    ctx->atol    = TOL;
+
+    ctx->logger.log = tdi_log_default_callback;
+
+    ret = tdi_nlsolve_context_init(ctx);
+    if (ret < 0)
+        goto finish;
+
+    // seed the initial guess with garbage, but not too large, numbers
+    for (int i = 0; i < N_X; i++)
+        for (int j = 0; j < N_Z; j++)
+            coeffs[j * N_X + i] = 0.2 * j + 0.1 * sin(138.0 * i);
+
+    ret = tdi_nlsolve_solve(ctx, scalar_linear_func, NULL, basis, coeffs);
+    if (ret < 0)
+        goto finish;
+
+    for (int j = 0; j < N_Z; j++)
+        for (int i = 0; i < N_X; i++) {
+            const int idx = j * N_X + i;
+            const double val = (i == FUNC_X && j == FUNC_Z) ? 1.0 : 0.0;
+
+            if (fabs(coeffs[idx] - val) > TOL) {
+                fprintf(stderr, "unexpected value %g at %d/%d\n",
+                        coeffs[idx], i, j);
+                ret = -1;
+                goto finish;
+            }
+        }
+
+finish:
+    tdi_nlsolve_context_free(&ctx);
+    tdi_basis_free(&basis);
+
+    return ret;
+}
+#undef N_X
+#undef N_Z
+#undef FUNC_X
+#undef FUNC_Z
+#undef TOL
+#undef ITER
+
+#define N_X 16
+#define N_Z 16
+#define L   0.1
+#define FUNC_X 4
+#define FUNC_Z 2
+#define ATOL 1e-14
+#define TOL 5e-15
+#define ITER 10
+static int
+scalar_quadratic_func(void *opaque, unsigned int eq_idx,
+                      const unsigned int *colloc_grid_order,
+                      const double * const *colloc_grid,
+                      const double * const (*vars)[PSSOLVE_DIFF_ORDER_NB],
+                      double *dst)
+{
+    BasisSetContext *basis = opaque;
+
+    for (int j = 0; j < colloc_grid_order[1]; j++) {
+        const double z = colloc_grid[1][j];
+
+        for (int i = 0; i < colloc_grid_order[0]; i++) {
+            const double x = colloc_grid[0][i];
+
+            const int idx = j * colloc_grid_order[0] + i;
+
+            const double d2x = vars[0][PSSOLVE_DIFF_ORDER_20][idx];
+            const double d2z = vars[0][PSSOLVE_DIFF_ORDER_02][idx];
+            const double val = vars[0][PSSOLVE_DIFF_ORDER_00][idx];
+
+            const double d2x_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_DIFF2, x, FUNC_X) *
+                                  tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, z, FUNC_Z);
+            const double d2z_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, x, FUNC_X) *
+                                  tdi_basis_eval(basis, BS_EVAL_TYPE_DIFF2, z, FUNC_Z);
+            const double val_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, x, FUNC_X) *
+                                  tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, z, FUNC_Z);
+
+            dst[idx] = (d2x - d2x_ex) + (d2z - d2z_ex) + (SQR(val) - SQR(val_ex));
+        }
+    }
+
+    return 0;
+}
+
+static int
+scalar_quadratic_jac(void *opaque, unsigned int eq_idx,
+                     unsigned int var_idx, enum PSSolveDiffOrder diff_order,
+                     const unsigned int *colloc_grid_order,
+                     const double * const *colloc_grid,
+                     const double * const (*vars)[PSSOLVE_DIFF_ORDER_NB],
+                     double *dst)
+{
+    BasisSetContext *basis = opaque;
+
+    for (int j = 0; j < colloc_grid_order[1]; j++) {
+        const double z = colloc_grid[1][j];
+
+        for (int i = 0; i < colloc_grid_order[0]; i++) {
+            const double x = colloc_grid[0][i];
+
+            const int idx = j * colloc_grid_order[0] + i;
+
+            const double val = vars[0][PSSOLVE_DIFF_ORDER_00][idx];
+
+            if (diff_order == PSSOLVE_DIFF_ORDER_20 ||
+                diff_order == PSSOLVE_DIFF_ORDER_02) {
+                dst[idx] = 1.0;
+            } else if (diff_order == PSSOLVE_DIFF_ORDER_00) {
+                dst[idx] = 2.0 * val;
+            } else {
+                dst[idx] = 0.0;
+            }
+        }
+    }
+
+    return 0;
+}
+
+static int scalar_quadratic(void)
+{
+    BasisSetContext *basis = NULL;
+    NLSolveContext    *ctx = NULL;
+
+    double coeffs[N_X * N_Z];
+    int ret;
+
+    ret = tdi_basis_init(&basis, BASIS_FAMILY_SB_EVEN, L);
+    if (ret < 0)
+        goto finish;
+
+    ret = tdi_nlsolve_context_alloc(&ctx, 1);
+    if (ret < 0)
+        return ret;
+
+    ctx->basis[0][0] = basis;
+    ctx->basis[0][1] = basis;
+
+    ctx->solve_order[0][0] = N_X;
+    ctx->solve_order[0][1] = N_Z;
+
+    ctx->maxiter = ITER;
+    ctx->atol    = ATOL;
+
+    ctx->logger.log = tdi_log_default_callback;
+
+    ret = tdi_nlsolve_context_init(ctx);
+    if (ret < 0)
+        goto finish;
+
+    memset(coeffs, 0, sizeof(coeffs));
+
+    ret = tdi_nlsolve_solve(ctx, scalar_quadratic_func, scalar_quadratic_jac, basis, coeffs);
+    if (ret < 0)
+        goto finish;
+
+    for (int j = 0; j < N_Z; j++)
+        for (int i = 0; i < N_X; i++) {
+            const int idx = j * N_X + i;
+            const double val = (i == FUNC_X && j == FUNC_Z) ? 1.0 : 0.0;
+
+            if (fabs(coeffs[idx] - val) > TOL) {
+                fprintf(stderr, "unexpected value %g at %d/%d\n",
+                        coeffs[idx], i, j);
+                ret = -1;
+                goto finish;
+            }
+        }
+
+finish:
+    tdi_nlsolve_context_free(&ctx);
+    tdi_basis_free(&basis);
+
+    return ret;
+}
+#undef N_X
+#undef N_Z
+#undef FUNC_X
+#undef FUNC_Z
+#undef TOL
+#undef ITER
+
+static const struct {
+    const char *name;
+    int (*test)(void);
+} tests[] = {
+    { "scalar_linear",    scalar_linear },
+    { "scalar_quadratic", scalar_quadratic },
+};
+
+int main(void)
+{
+    int ret = 0;
+
+    for (int i = 0; i < ARRAY_ELEMS(tests); i++) {
+        fprintf(stderr, "executing test '%s'\n", tests[i].name);
+        ret = tests[i].test();
+        if (ret < 0) {
+            fprintf(stderr, "test '%s' failed\n", tests[i].name);
+            return -1;
+        }
+        fprintf(stderr, "test '%s' succeeded\n", tests[i].name);
+    }
+
+    return 0;
+}