eval: refactor evaluating the metric

1 files changed, 43 insertions, 137 deletions

diff --git a/eval.c b/eval.c
index c8b4b95..e724447 100644
--- a/eval.c
+++ b/eval.c
@@ -105,103 +105,6 @@ fail:
 
 }
 
-static void eval_metric(const BDContext *bd,
-                        const double *rho, int nb_coords_rho,
-                        const double *z,   int nb_coords_z,
-                        enum BDMetricComponent comp,
-                        const unsigned int diff_order[2],
-                        const double *psi,
-                        const double *dpsi_rho, const double *dpsi_z,
-                        const double *d2psi_rho, const double *d2psi_z, const double *d2psi_rho_z,
-                        double *out, unsigned int out_stride)
-{
-    const BDPriv *s = bd->priv;
-
-    if (comp != BD_METRIC_COMPONENT_RHORHO &&
-        comp != BD_METRIC_COMPONENT_ZZ     &&
-        comp != BD_METRIC_COMPONENT_PHIPHI) {
-        memset(out, 0, out_stride * nb_coords_z * sizeof(*out));
-        return;
-    }
-
-/* the template for the loop over the grid points */
-#define EVAL_LOOP(DO_EVAL)                                          \
-do {                                                                \
-    for (int i = 0; i < nb_coords_z; i++) {                         \
-        const double zz = z[i];                                     \
-        double *dst = out + i * out_stride;                         \
-                                                                    \
-        for (int j = 0; j < nb_coords_rho; j++) {                   \
-            const double rrho = rho[j];                             \
-            const double ppsi = psi[i * nb_coords_rho + j];         \
-            const double psi2 = SQR(ppsi);                          \
-            const double psi3 = psi2 * ppsi;                        \
-            const double psi4 = SQR(psi2);                          \
-            double val;                                             \
-                                                                    \
-            DO_EVAL;                                                \
-                                                                    \
-            *dst++ = val;                                           \
-        }                                                           \
-    }                                                               \
-} while (0)
-
-#define GRID(arr) (arr[i * nb_coords_rho + j])
-
-    if (comp == BD_METRIC_COMPONENT_PHIPHI) {
-        switch (diff_order[0]) {
-        case 0:
-            switch (diff_order[1]) {
-            case 0: EVAL_LOOP(val = SQR(rrho) * psi4);                                                                                                                    break;
-            case 1: EVAL_LOOP(val = 4 * SQR(rrho) * psi3 * GRID(dpsi_z));                                                                                                 break; // ∂_z
-            case 2: EVAL_LOOP(val = 4 * SQR(rrho) * (GRID(d2psi_z) * psi3 + 3 * psi2 * SQR(GRID(dpsi_z))));                                                               break; // ∂_z ∂_z
-            }
-            break;
-        case 1:
-            switch (diff_order[1]) {
-            case 0: EVAL_LOOP(val = 2 * rrho * psi4 + 4 * SQR(rrho) * psi3 * GRID(dpsi_rho));                                                                             break; // ∂_ρ
-            case 1: EVAL_LOOP(val = 8 * rrho * psi3 * GRID(dpsi_z) + 12 * SQR(rrho) * psi2 * GRID(dpsi_z) * GRID(dpsi_rho) + 4 * SQR(rrho) * psi3 * GRID(d2psi_rho_z));   break; // ∂_ρ ∂_z
-            }
-            break;
-        case 2:     EVAL_LOOP(val = 4 * SQR(rrho) * psi3 * GRID(d2psi_rho) + 12 * SQR(rrho * ppsi * GRID(dpsi_rho)) + 16 * rrho * psi3 * GRID(dpsi_rho) + 2 * psi4);      break; // ∂_ρ ∂_ρ
-        }
-    } else {
-        // γ_ρρ / γ_zz
-
-/* a wrapper around the actual evaluation expression that provides the q function and its
- * derivatives as needed */
-#define DO_EVAL_Q_WRAPPER(DO_EVAL, eval_drho, eval_dz, eval_d2rho, eval_d2z, eval_d2rhoz)                                                    \
-do {                                                                                                                                         \
-    const double base = psi4 * exp(2 * bdi_qfunc_eval(s->qfunc, rrho, zz, 0));                                                               \
-    double drho, dz, d2rho, d2z, d2rho_z;                                                                                                    \
-                                                                                                                                             \
-    if (eval_drho)   drho    = bdi_qfunc_eval(s->qfunc, rrho, zz, 1) + 2 * GRID(dpsi_rho)    / ppsi;                                            \
-    if (eval_dz)     dz      = bdi_qfunc_eval(s->qfunc, rrho, zz, 3) + 2 * GRID(dpsi_z)      / ppsi;                                            \
-    if (eval_d2rho)  d2rho   = bdi_qfunc_eval(s->qfunc, rrho, zz, 2) + 2 * GRID(d2psi_rho)   / ppsi - 2 * SQR(GRID(dpsi_rho) / ppsi);           \
-    if (eval_d2z)    d2z     = bdi_qfunc_eval(s->qfunc, rrho, zz, 6) + 2 * GRID(d2psi_z)     / ppsi - 2 * SQR(GRID(dpsi_z)   / ppsi);           \
-    if (eval_d2rhoz) d2rho_z = bdi_qfunc_eval(s->qfunc, rrho, zz, 4) + 2 * GRID(d2psi_rho_z) / ppsi - 2 * GRID(dpsi_rho) * GRID(dpsi_z) / psi2; \
-                                                                                                                                             \
-    DO_EVAL;                                                                                                                                 \
-} while (0)
-
-        switch (diff_order[0]) {
-        case 0:
-            switch (diff_order[1]) {
-            case 0: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = base,                                      0, 0, 0, 0, 0)); break;
-            case 1: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 2 * base * dz,                             0, 1, 0, 0, 0)); break;        // ∂_z
-            case 2: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 4 * SQR(dz) * base + 2 * base * d2z,       0, 1, 0, 1, 0)); break;        // ∂_z ∂_z
-            }
-            break;
-        case 1:
-            switch (diff_order[1]) {
-            case 0: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 2 * base * drho,                           1, 0, 0, 0, 0)); break;        // ∂_ρ
-            case 1: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 4 * base * drho * dz + 2 * base * d2rho_z, 1, 1, 0, 0, 1)); break;        // ∂_ρ ∂_z
-            }
-            break;
-        case 2:     EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 4 * SQR(drho) * base + 2 * base * d2rho,   1, 0, 1, 0, 0)); break;        // ∂_ρ ∂_ρ
-        }
-    }
-}
 
 int bd_eval_metric(const BDContext *bd, const double *rho, int nb_coords_rho,
                    const double *z, int nb_coords_z,
@@ -210,7 +113,7 @@ int bd_eval_metric(const BDContext *bd, const double *rho, int nb_coords_rho,
                    double **out, unsigned int *out_strides)
 {
     const int nb_coords = nb_coords_rho * nb_coords_z;
-    double *psi = NULL, *dpsi_rho = NULL, *dpsi_z = NULL, *d2psi_rho = NULL, *d2psi_rho_z = NULL, *d2psi_z = NULL;
+    double *psi[9] = { NULL }, *q[9] = { NULL };
     int ret = 0;
 
     /* check the parameters for validity */
@@ -227,56 +130,59 @@ int bd_eval_metric(const BDContext *bd, const double *rho, int nb_coords_rho,
     }
 
     /* evaluate the conformal factor and its derivatives as necessary */
-#define EVAL_PSI(arr, order)                                    \
-do {                                                            \
-    if (arr)                                                    \
-        break;                                                  \
-                                                                \
-    arr = malloc(nb_coords * sizeof(*arr));                     \
-    if (!arr) {                                                 \
-        ret = -ENOMEM;                                          \
-        goto fail;                                              \
-    }                                                           \
-                                                                \
-    ret = bd_eval_psi(bd, rho, nb_coords_rho, z, nb_coords_z,   \
-                      order, arr, nb_coords_rho);               \
-    if (ret < 0)                                                \
-        goto fail;                                              \
-} while (0)
-
     for (int i = 0; i < nb_comp; i++) {
         if (comp[i] != BD_METRIC_COMPONENT_RHORHO &&
             comp[i] != BD_METRIC_COMPONENT_ZZ     &&
             comp[i] != BD_METRIC_COMPONENT_PHIPHI)
             continue;
 
-        EVAL_PSI(psi,             ((unsigned int[2]){ 0, 0 }));
-        if (diff_order[i][0])
-            EVAL_PSI(dpsi_rho,    ((unsigned int[2]){ 1, 0 }));
-        if (diff_order[i][0] > 1)
-            EVAL_PSI(d2psi_rho,   ((unsigned int[2]){ 2, 0 }));
-        if (diff_order[i][1])
-            EVAL_PSI(dpsi_z,      ((unsigned int[2]){ 0, 1 }));
-        if (diff_order[i][1] > 1)
-            EVAL_PSI(d2psi_z,     ((unsigned int[2]){ 0, 2 }));
-        if (diff_order[i][0] && diff_order[i][1])
-            EVAL_PSI(d2psi_rho_z, ((unsigned int[2]){ 1, 1 }));
+        for (int j = 0; j < ARRAY_ELEMS(psi); j++) {
+            if (psi[j])
+                continue;
+            if (j / 3 + j % 3 > 2)
+                continue;
+            if (!(j % 3 <= diff_order[i][0] ||
+                  j / 3 <= diff_order[i][1]))
+                continue;
+
+            psi[j] = malloc(nb_coords * sizeof(*psi[j]));
+            if (!psi[j]) {
+                ret = -ENOMEM;
+                goto fail;
+            }
+
+            ret = bd_eval_psi(bd, rho, nb_coords_rho, z, nb_coords_z,
+                              (unsigned int[2]){ j % 3, j / 3 }, psi[j],
+                              nb_coords_rho);
+            if (ret < 0)
+                goto fail;
+
+            q[j] = malloc(nb_coords * sizeof(*q[j]));
+            if (!q[j]) {
+                ret = -ENOMEM;
+                goto fail;
+            }
+
+            ret = bd_eval_q(bd, rho, nb_coords_rho, z, nb_coords_z,
+                            (unsigned int[2]){ j % 3, j / 3 }, q[j],
+                            nb_coords_rho);
+            if (ret < 0)
+                goto fail;
+        }
     }
 
     /* evaluate the requested metric components */
-    for (int i = 0; i < nb_comp; i++)
-        eval_metric(bd, rho, nb_coords_rho, z, nb_coords_z,
-                    comp[i], diff_order[i],
-                    psi, dpsi_rho, dpsi_z, d2psi_rho, d2psi_z, d2psi_rho_z,
-                    out[i], out_strides[i]);
+    for (int i = 0; i < nb_comp; i++) {
+        bdi_eval_metric(bd, rho, nb_coords_rho, z, nb_coords_z,
+                        comp[i], diff_order[i],
+                        psi, q, out[i], out_strides[i]);
+    }
 
 fail:
-    free(psi);
-    free(dpsi_rho);
-    free(dpsi_z);
-    free(d2psi_rho);
-    free(d2psi_rho_z);
-    free(d2psi_z);
+    for (int i = 0; i < ARRAY_ELEMS(psi); i++)
+        free(psi[i]);
+    for (int i = 0; i < ARRAY_ELEMS(q); i++)
+        free(q[i]);
 
     return ret;
 }