tests/nlsolve: add a vector quadratic test

1 files changed, 142 insertions, 1 deletions

diff --git a/tests/nlsolve.c b/tests/nlsolve.c
index 6d46c9e..502164a 100644
--- a/tests/nlsolve.c
+++ b/tests/nlsolve.c
@@ -140,7 +140,7 @@ finish:
 #define FUNC_X 4
 #define FUNC_Z 2
 #define ATOL 1e-14
-#define TOL 5e-15
+#define TOL 1e-14
 #define ITER 10
 static int
 scalar_quadratic_func(void *opaque, unsigned int eq_idx,
@@ -269,8 +269,148 @@ finish:
 }
 #undef N_X
 #undef N_Z
+#undef L
 #undef FUNC_X
 #undef FUNC_Z
+#undef ATOL
+#undef TOL
+#undef ITER
+
+#define N_R   16
+#define N_PHI 8
+#define L   0.5
+#define FUNC0_R   4
+#define FUNC0_PHI 2
+#define FUNC1_R   2
+#define FUNC1_PHI 3
+#define ATOL 1e-14
+#define TOL 1e-14
+#define ITER 128
+static int
+vector_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;
+    unsigned int exact_r   = eq_idx == 0 ? FUNC0_R : FUNC1_R;
+    unsigned int exact_phi = eq_idx == 0 ? FUNC0_PHI : FUNC1_PHI;
+
+    for (int j = 0; j < colloc_grid_order[1]; j++) {
+        const double phi = colloc_grid[1][j];
+
+        for (int i = 0; i < colloc_grid_order[0]; i++) {
+            const double r = colloc_grid[0][i];
+
+            const int idx = j * colloc_grid_order[0] + i;
+
+            const double  dr   = vars[eq_idx][PSSOLVE_DIFF_ORDER_10][idx];
+            const double d2r   = vars[eq_idx][PSSOLVE_DIFF_ORDER_20][idx];
+            const double d2phi = vars[eq_idx][PSSOLVE_DIFF_ORDER_02][idx];
+            const double val0 = vars[0][PSSOLVE_DIFF_ORDER_00][idx];
+            const double val1 = vars[1][PSSOLVE_DIFF_ORDER_00][idx];
+
+            const double  dr_ex  = tdi_basis_eval(basis[0], BS_EVAL_TYPE_DIFF1, r,    exact_r) *
+                                   tdi_basis_eval(basis[1], BS_EVAL_TYPE_VALUE, phi,  exact_phi);
+            const double d2r_ex  = tdi_basis_eval(basis[0], BS_EVAL_TYPE_DIFF2, r,    exact_r) *
+                                   tdi_basis_eval(basis[1], BS_EVAL_TYPE_VALUE, phi,  exact_phi);
+            const double d2phi_ex = tdi_basis_eval(basis[0], BS_EVAL_TYPE_VALUE, r,   exact_r) *
+                                    tdi_basis_eval(basis[1], BS_EVAL_TYPE_DIFF2, phi, exact_phi);
+            const double val0_ex = tdi_basis_eval(basis[0], BS_EVAL_TYPE_VALUE, r,   FUNC0_R) *
+                                   tdi_basis_eval(basis[1], BS_EVAL_TYPE_VALUE, phi, FUNC0_PHI);
+            const double val1_ex = tdi_basis_eval(basis[0], BS_EVAL_TYPE_VALUE, r,   FUNC1_R) *
+                                   tdi_basis_eval(basis[1], BS_EVAL_TYPE_VALUE, phi, FUNC1_PHI);
+
+            dst[idx] = (d2r - d2r_ex) + (dr - dr_ex) / r + (d2phi - d2phi_ex) / SQR(r) + (val0 * val1 - val0_ex * val1_ex);
+        }
+    }
+
+    return 0;
+}
+
+static int vector_quadratic(void)
+{
+    BasisSetContext *basis[2] = { NULL };
+    NLSolveContext    *ctx = NULL;
+
+    double coeffs[2 * N_R * N_PHI];
+    int ret;
+
+    ret = tdi_basis_init(&basis[0], BASIS_FAMILY_SB_EVEN,  L);
+    ret = tdi_basis_init(&basis[1], BASIS_FAMILY_COS_EVEN, L);
+    if (ret < 0)
+        goto finish;
+
+    ret = tdi_nlsolve_context_alloc(&ctx, 2);
+    if (ret < 0)
+        return ret;
+
+    ctx->basis[0][0] = basis[0];
+    ctx->basis[1][0] = basis[0];
+    ctx->basis[0][1] = basis[1];
+    ctx->basis[1][1] = basis[1];
+
+    ctx->solve_order[0][0] = N_R;
+    ctx->solve_order[0][1] = N_PHI;
+    ctx->solve_order[1][0] = N_R;
+    ctx->solve_order[1][1] = N_PHI;
+
+    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, vector_quadratic_func, NULL, basis, coeffs);
+    if (ret < 0)
+        goto finish;
+
+    for (int j = 0; j < N_PHI; j++)
+        for (int i = 0; i < N_R; i++) {
+            const int idx = j * N_R + i;
+            const double val = (i == FUNC0_R && j == FUNC0_PHI) ? 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;
+            }
+        }
+    for (int j = 0; j < N_PHI; j++)
+        for (int i = 0; i < N_R; i++) {
+            const int idx = N_PHI * N_R + j * N_R + i;
+            const double val = (i == FUNC1_R && j == FUNC1_PHI) ? 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[0]);
+    tdi_basis_free(&basis[1]);
+
+    return ret;
+}
+#undef N_R
+#undef N_PHI
+#undef L
+#undef FUNC0_R
+#undef FUNC0_PHI
+#undef FUNC1_R
+#undef FUNC1_PHI
+#undef ATOL
 #undef TOL
 #undef ITER
 
@@ -280,6 +420,7 @@ static const struct {
 } tests[] = {
     { "scalar_linear",    scalar_linear },
     { "scalar_quadratic", scalar_quadratic },
+    { "vector_quadratic", vector_quadratic },
 };
 
 int main(void)