Switch to polar coordinates in the solver.

1 files changed, 91 insertions, 2 deletions

diff --git a/solve.c b/solve.c
index fe47832..fa527a9 100644
--- a/solve.c
+++ b/solve.c
@@ -34,7 +34,7 @@
 #include "internal.h"
 #include "qfunc.h"
 
-static int brill_construct_matrix(const BDContext *bd, double *mat,
+static int brill_construct_matrix_cylindrical(const BDContext *bd, double *mat,
                                   double *rhs)
 {
     BDPriv *s = bd->priv;
@@ -104,6 +104,90 @@ fail:
     return ret;
 }
 
+static int brill_construct_matrix_polar(const BDContext *bd, double *mat,
+                                        double *rhs)
+{
+    BDPriv *s = bd->priv;
+
+    double *basis_val, *basis_dval, *basis_d2val;
+    int idx_coeff_0, idx_coeff_1, idx_grid_0, idx_grid_1;
+
+    double *basis[2][3] = { { NULL } };
+    int ret = 0;
+
+    /* precompute the basis values we will need */
+    for (int i = 0; i < ARRAY_ELEMS(basis); i++) {
+        for (int j = 0; j < ARRAY_ELEMS(basis[i]); j++) {
+            basis[i][j] = malloc(sizeof(*basis[i][j]) * s->nb_colloc_points[i] * s->nb_coeffs[i]);
+            if (!basis[i][j]) {
+                ret = -ENOMEM;
+                goto fail;
+            }
+        }
+        for (int j = 0; j < s->nb_colloc_points[i]; j++) {
+            double coord = bdi_basis_colloc_point(s->basis[i], j, s->nb_coeffs[i]);
+            for (int k = 0; k < s->nb_coeffs[i]; k++) {
+                basis[i][0][j * s->nb_coeffs[i] + k] = bdi_basis_eval(s->basis[i], BS_EVAL_TYPE_VALUE, coord, k);
+                basis[i][1][j * s->nb_coeffs[i] + k] = bdi_basis_eval(s->basis[i], BS_EVAL_TYPE_DIFF1, coord, k);
+                basis[i][2][j * s->nb_coeffs[i] + k] = bdi_basis_eval(s->basis[i], BS_EVAL_TYPE_DIFF2, coord, k);
+            }
+        }
+    }
+
+#define   BASIS0 (basis[0][0][idx_grid_0 * s->nb_coeffs[0] + idx_coeff_0])
+#define  DBASIS0 (basis[0][1][idx_grid_0 * s->nb_coeffs[0] + idx_coeff_0])
+#define D2BASIS0 (basis[0][2][idx_grid_0 * s->nb_coeffs[0] + idx_coeff_0])
+#define   BASIS1 (basis[1][0][idx_grid_1 * s->nb_coeffs[1] + idx_coeff_1])
+#define  DBASIS1 (basis[1][1][idx_grid_1 * s->nb_coeffs[1] + idx_coeff_1])
+#define D2BASIS1 (basis[1][2][idx_grid_1 * s->nb_coeffs[1] + idx_coeff_1])
+
+    for (idx_grid_1 = 0; idx_grid_1 < s->nb_colloc_points[1]; idx_grid_1++) {
+        double phi = bdi_basis_colloc_point(s->basis[1], idx_grid_1, s->nb_coeffs[1]);
+
+        for (idx_grid_0 = 0; idx_grid_0 < s->nb_colloc_points[0]; idx_grid_0++) {
+            double r = bdi_basis_colloc_point(s->basis[0], idx_grid_0, s->nb_coeffs[0]);
+
+            double x = r * cos(phi);
+            double z = r * sin(phi);
+
+            double d2q = bdi_qfunc_eval(s->qfunc, x, z, 2) +
+                         bdi_qfunc_eval(s->qfunc, x, z, 6);
+            int idx_grid  = idx_grid_1 * s->nb_colloc_points[0] + idx_grid_0;
+
+            for (idx_coeff_1 = 0; idx_coeff_1 < s->nb_coeffs[1]; idx_coeff_1++)
+                for (idx_coeff_0 = 0; idx_coeff_0 < s->nb_coeffs[0]; idx_coeff_0++) {
+                    int idx_coeff = idx_coeff_1 * s->nb_coeffs[0] + idx_coeff_0;
+
+                    double basis_val_20 =  ((r > EPS) ? (SQR(x / r) * D2BASIS0 * BASIS1 + SQR(z / SQR(r)) * BASIS0 * D2BASIS1
+                                                         + (1.0 - SQR(x / r)) / r * DBASIS0 * BASIS1
+                                                         + 2 * x * z / SQR(SQR(r)) * BASIS0 * DBASIS1
+                                                         - 2 * z * x / (r * SQR(r)) * DBASIS0 * DBASIS1) : 0.0);
+                    double basis_val_02 =  ((r > EPS) ? (SQR(z / r) * D2BASIS0 * BASIS1 + SQR(x / SQR(r)) * BASIS0 * D2BASIS1
+                                                           + (1.0 - SQR(z / r)) / r * DBASIS0 * BASIS1
+                                                           - 2 * x * z / SQR(SQR(r)) * BASIS0 * DBASIS1
+                                                           + 2 * z * x / (r * SQR(r)) * DBASIS0 * DBASIS1) : 0.0);
+                    double basis_val_10 = ((r > EPS) ? (DBASIS0 * BASIS1 * x / r - BASIS0 * DBASIS1 * z / SQR(r)) : 0.0);
+
+                    double val  = basis_val_20 + basis_val_02 + BASIS0 * BASIS1 * 0.25 * d2q;
+                    if (fabs(x) > EPS)
+                        val += basis_val_10 / fabs(x);
+                    else
+                        val += basis_val_20;
+
+                    mat[idx_grid + NB_COLLOC_POINTS(s) * idx_coeff] = val;
+                }
+            rhs[idx_grid] = -0.25 * d2q;
+        }
+    }
+
+fail:
+    for (int i = 0; i < ARRAY_ELEMS(basis); i++)
+        for (int j = 0; j < ARRAY_ELEMS(basis[i]); j++)
+            free(basis[i][j]);
+
+    return ret;
+}
+
 static int brill_solve_linear(const BDContext *bd, double *mat, double **px, double **prhs)
 {
     const BDPriv *s = bd->priv;
@@ -223,7 +307,12 @@ int bdi_solve(BDContext *bd)
     }
 
     /* fill the matrix */
-    ret = brill_construct_matrix(bd, mat, rhs);
+    switch (s->coord_system) {
+    case COORD_SYSTEM_CYLINDRICAL: ret = brill_construct_matrix_cylindrical(bd, mat, rhs);
+                                   break;
+    case COORD_SYSTEM_RADIAL:      ret = brill_construct_matrix_polar(bd, mat, rhs);
+                                   break;
+    }
     if (ret < 0)
         goto fail;