Switch to polar coordinates in the solver.

1 files changed, 147 insertions, 11 deletions

diff --git a/eval.c b/eval.c
index e724447..a27d0f5 100644
--- a/eval.c
+++ b/eval.c
@@ -29,30 +29,140 @@
 #include "internal.h"
 #include "qfunc.h"
 
-int bd_eval_psi(const BDContext *bd,
-                const double *rho, int nb_coords_rho,
-                const double *z,   int nb_coords_z,
-                const unsigned int diff_order[2],
-                double *psi, unsigned int psi_stride)
+static int eval_psi_radial(const BDContext *bd,
+                           const double *rho, int nb_coords_rho,
+                           const double *z,   int nb_coords_z,
+                           const unsigned int diff_order[2],
+                           double *psi, unsigned int psi_stride)
 {
     BDPriv *s = bd->priv;
     enum BSEvalType type;
 
-    double *basis_val_rho = NULL, *basis_val_z = NULL;
+    double *basis_val[2][3] = { { NULL } };
 
     double add = 0.0;
 
     int ret = 0;
 
-    if (diff_order[0] > 2 || diff_order[1] > 2) {
-        bdi_log(bd, 0, "Derivatives of higher order than 2 are not supported.\n");
-        return -ENOSYS;
-    }
-
     if (diff_order[0] == 0 && diff_order[1] == 0)
         add = 1.0;
 
     /* precompute the basis values on the grid points */
+    for (int i = 0; i < ARRAY_ELEMS(basis_val[0]); i++) {
+        posix_memalign((void**)&basis_val[0][i], REQ_ALIGNMENT(char), sizeof(*basis_val[0][i]) * s->nb_coeffs[0]);
+        posix_memalign((void**)&basis_val[1][i], REQ_ALIGNMENT(char), sizeof(*basis_val[1][i]) * s->nb_coeffs[1]);
+        if (!basis_val[0][i] || !basis_val[1][i]) {
+            ret = -ENOMEM;
+            goto fail;
+        }
+    }
+
+#define   BASIS0 (basis_val[0][0][n])
+#define  DBASIS0 (basis_val[0][1][n])
+#define D2BASIS0 (basis_val[0][2][n])
+#define   BASIS1 (basis_val[1][0][m])
+#define  DBASIS1 (basis_val[1][1][m])
+#define D2BASIS1 (basis_val[1][2][m])
+
+    for (int i = 0; i < nb_coords_z; i++) {
+        double *dst = psi + i * psi_stride;
+        double zz = z[i];
+
+        for (int j = 0; j < nb_coords_rho; j++) {
+            double ppsi = add;
+            double rrho = rho[j];
+
+            double r    = sqrt(SQR(rrho) + SQR(zz));
+            double phi  = atan2(zz, rrho);
+
+            bdi_basis_eval_multi(s->basis[0], r,   basis_val[0], 0, s->nb_coeffs[0]);
+            bdi_basis_eval_multi(s->basis[1], phi, basis_val[1], 0, s->nb_coeffs[1]);
+
+            if (diff_order[0] == 0 && diff_order[1] == 0) {
+                for (int m = 0; m < s->nb_coeffs[1]; m++) {
+                    double tmp = 0.0;
+                    for (int n = 0; n < s->nb_coeffs[0]; n++) {
+                        tmp += s->coeffs[m * s->coeffs_stride + n] * BASIS0;
+                    }
+                    ppsi += tmp * BASIS1;
+                }
+            } else if (diff_order[0] == 2 && diff_order[1] == 0) {
+                for (int m = 0; m < s->nb_coeffs[1]; m++) {
+                    for (int n = 0; n < s->nb_coeffs[0]; n++) {
+                        double basis_val_20 =  ((r > EPS) ? (SQR(rrho / r) * D2BASIS0 * BASIS1 + SQR(zz / SQR(r)) * BASIS0 * D2BASIS1
+                                                             + (1.0 - SQR(rrho / r)) / r * DBASIS0 * BASIS1
+                                                             + 2 * rrho * zz / SQR(SQR(r)) * BASIS0 * DBASIS1
+                                                             - 2 * zz * rrho / (r * SQR(r)) * DBASIS0 * DBASIS1) : 0.0);
+                        ppsi += s->coeffs[m * s->coeffs_stride + n] * basis_val_20;
+                    }
+                }
+            } else if (diff_order[0] == 0 && diff_order[1] == 2) {
+                for (int m = 0; m < s->nb_coeffs[1]; m++) {
+                    for (int n = 0; n < s->nb_coeffs[0]; n++) {
+                        double basis_val_02 =  ((r > EPS) ? (SQR(zz / r) * D2BASIS0 * BASIS1 + SQR(rrho / SQR(r)) * BASIS0 * D2BASIS1
+                                                               + (1.0 - SQR(zz / r)) / r * DBASIS0 * BASIS1
+                                                               - 2 * rrho * zz / SQR(SQR(r)) * BASIS0 * DBASIS1
+                                                               + 2 * zz * rrho / (r * SQR(r)) * DBASIS0 * DBASIS1) : 0.0);
+                        ppsi += s->coeffs[m * s->coeffs_stride + n] * basis_val_02;
+                    }
+                }
+            } else if (diff_order[0] == 1 && diff_order[1] == 0) {
+                for (int m = 0; m < s->nb_coeffs[1]; m++) {
+                    for (int n = 0; n < s->nb_coeffs[0]; n++) {
+                        double basis_val_10 = ((r > EPS) ? (DBASIS0 * BASIS1 * rrho / r - BASIS0 * DBASIS1 * zz / SQR(r)) : 0.0);
+
+                        ppsi += s->coeffs[m * s->coeffs_stride + n] * basis_val_10;
+                    }
+                }
+            } else if (diff_order[0] == 0 && diff_order[1] == 1) {
+                for (int m = 0; m < s->nb_coeffs[1]; m++) {
+                    for (int n = 0; n < s->nb_coeffs[0]; n++) {
+                        double basis_val_01 = ((r > EPS) ? (DBASIS0 * BASIS1 * zz / r + BASIS0 * DBASIS1 * rrho / SQR(r)) : 0.0);
+
+                        ppsi += s->coeffs[m * s->coeffs_stride + n] * basis_val_01;
+                    }
+                }
+            } else if (diff_order[0] == 1 && diff_order[1] == 1) {
+                for (int m = 0; m < s->nb_coeffs[1]; m++) {
+                    for (int n = 0; n < s->nb_coeffs[0]; n++) {
+                        double basis_val_11 = ((r > EPS) ? (rrho * zz / SQR(r) * D2BASIS0 * BASIS1 - rrho * zz / SQR(SQR(r)) * BASIS0 * D2BASIS1
+                                                            - rrho * zz / (r * SQR(r)) * DBASIS0 * BASIS1
+                                                            - (1.0 - SQR(zz / r)) / SQR(r) * BASIS0 * DBASIS1
+                                                            + (SQR(rrho) - SQR(zz)) / (r * SQR(r)) * DBASIS0 * DBASIS1) : 0.0);
+
+                        ppsi += s->coeffs[m * s->coeffs_stride + n] * basis_val_11;
+                    }
+                }
+            }
+            *dst++ = ppsi;
+        }
+    }
+
+fail:
+    for (int i = 0; i < ARRAY_ELEMS(basis_val[0]); i++) {
+        free(basis_val[0][i]);
+        free(basis_val[1][i]);
+    }
+
+    return ret;
+}
+
+static int eval_psi_cylindric(const BDContext *bd,
+                              const double *rho, int nb_coords_rho,
+                              const double *z,   int nb_coords_z,
+                              const unsigned int diff_order[2],
+                              double *psi, unsigned int psi_stride)
+{
+    BDPriv *s = bd->priv;
+    enum BSEvalType type;
+
+    double *basis_val_rho = NULL, *basis_val_z = NULL;
+
+    double add = 0.0;
+
+    int ret = 0;
+
+    /* precompute the basis values on the grid points */
     basis_val_rho = malloc(sizeof(*basis_val_rho) * s->nb_coeffs[0] * nb_coords_rho);
     basis_val_z   = malloc(sizeof(*basis_val_z)   * s->nb_coeffs[1] * nb_coords_z);
     if (!basis_val_rho || !basis_val_z) {
@@ -105,6 +215,32 @@ fail:
 
 }
 
+int bd_eval_psi(const BDContext *bd,
+                const double *rho, int nb_coords_rho,
+                const double *z,   int nb_coords_z,
+                const unsigned int diff_order[2],
+                double *psi, unsigned int psi_stride)
+{
+    BDPriv *s = bd->priv;
+    int ret = 0;
+
+    if (diff_order[0] > 2 || diff_order[1] > 2) {
+        bdi_log(bd, 0, "Derivatives of higher order than 2 are not supported.\n");
+        return -ENOSYS;
+    }
+
+    switch (s->coord_system) {
+    case COORD_SYSTEM_CYLINDRICAL: ret = eval_psi_cylindric(bd, rho, nb_coords_rho, z, nb_coords_z,
+                                                            diff_order, psi, psi_stride);
+                                   break;
+    case COORD_SYSTEM_RADIAL:      ret = eval_psi_radial(bd, rho, nb_coords_rho, z, nb_coords_z,
+                                                         diff_order, psi, psi_stride);
+                                   break;
+    }
+
+    return ret;
+
+}
 
 int bd_eval_metric(const BDContext *bd, const double *rho, int nb_coords_rho,
                    const double *z, int nb_coords_z,