eval: refactor evaluating the metric

5 files changed, 290 insertions, 145 deletions

diff --git a/Makefile b/Makefile
index 0efa6c0..c868322 100644
--- a/Makefile
+++ b/Makefile
@@ -3,17 +3,19 @@ LDFLAGS = --version-script=libbrilldata.v -shared -lm -llapacke
 
 TARGET = libbrilldata.so
 
-OBJECTS = basis.o               \
-          eval.o                \
-          init.o                \
-          log.o                 \
-          qfunc.o               \
-          solve.o
+OBJS = basis.o               \
+       eval.o                \
+       eval_metric.o         \
+       init.o                \
+       log.o                 \
+       qfunc.o               \
+       solve.o               \
+
 
 all: $(TARGET)
 
-$(TARGET): $(OBJECTS)
-	ld ${LDFLAGS} -o $@ $(OBJECTS)
+$(TARGET): $(OBJS)
+	ld ${LDFLAGS} -o $@ $(OBJS)
 
 clean:
 	-rm $(OBJECTS) $(TARGET)
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;
 }
diff --git a/eval_metric.c b/eval_metric.c
new file mode 100644
index 0000000..f499e32
--- /dev/null
+++ b/eval_metric.c
@@ -0,0 +1,161 @@
+/*
+ * Copyright 2016 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 <math.h>
+#include <stddef.h>
+#include <string.h>
+
+#include "brill_data.h"
+#include "internal.h"
+
+#define GRID(arr) (arr[i * nb_coords_rho + j])
+
+
+/* φφ components */
+
+#define COMP phiphi
+
+#define DIFF 0
+#define DO_EVAL val = SQR(rrho) * psi4
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#define DIFF 3
+#define DO_EVAL val = 4 * SQR(rrho) * psi3 * GRID(psi[3])
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#define DIFF 6
+#define DO_EVAL val = 4 * SQR(rrho) * (GRID(psi[6]) * psi3 + 3 * psi2 * SQR(GRID(psi[3])))
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#define DIFF 1
+#define DO_EVAL val = 2 * rrho * psi4 + 4 * SQR(rrho) * psi3 * GRID(psi[1])
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#define DIFF 4
+#define DO_EVAL val =  8 * rrho * psi3 * GRID(psi[3]) +\
+                      12 * SQR(rrho) * psi2 * GRID(psi[3]) * GRID(psi[1]) +\
+                       4 * SQR(rrho) * psi3 * GRID(psi[4])
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#define DIFF 2
+#define DO_EVAL val =  4 * SQR(rrho) * psi3 * GRID(psi[2]) +\
+                      12 * SQR(rrho * ppsi * GRID(psi[1])) +\
+                      16 * rrho * psi3 * GRID(psi[1]) + 2 * psi4
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#undef COMP
+
+/* ρρ components */
+#define COMP rhorho
+
+#define DIFF 0
+#define DO_EVAL val = base
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#define DIFF 3
+#define DO_EVAL val = 2 * base * dqz
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#define DIFF 6
+#define DO_EVAL val = 4 * SQR(dqz) * base + 2 * base * dq2z
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#define DIFF 1
+#define DO_EVAL val = 2 * base * dqrho
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#define DIFF 4
+#define DO_EVAL val = 4 * base * dqrho * dqz + 2 * base * dq2rho_z
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#define DIFF 2
+#define DO_EVAL val = 4 * SQR(dqrho) * base + 2 * base * dq2rho
+#include "eval_metric_template.c"
+#undef DIFF
+#undef DO_EVAL
+
+#undef COMP
+
+typedef void (*EvalMetricFunc)(const BDContext *bd, double *out, ptrdiff_t out_stride,
+                               int nb_coords_rho, const double *rho,
+                               int nb_coords_z,   const double *z,
+                               double *psi[9], double *q[9]);
+
+static const EvalMetricFunc eval_metric_phiphi[9] = {
+    [0] = eval_metric_phiphi_0,
+    [1] = eval_metric_phiphi_1,
+    [2] = eval_metric_phiphi_2,
+    [3] = eval_metric_phiphi_3,
+    [4] = eval_metric_phiphi_4,
+    [6] = eval_metric_phiphi_6,
+};
+
+static const EvalMetricFunc eval_metric_rhorho[9] = {
+    [0] = eval_metric_rhorho_0,
+    [1] = eval_metric_rhorho_1,
+    [2] = eval_metric_rhorho_2,
+    [3] = eval_metric_rhorho_3,
+    [4] = eval_metric_rhorho_4,
+    [6] = eval_metric_rhorho_6,
+};
+
+void bdi_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],
+                     double *psi[9], double *q[9],
+                     double *out, ptrdiff_t out_stride)
+{
+    int diff_idx = diff_order[0] * 3 + diff_order[1];
+    const EvalMetricFunc *functab;
+
+    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;
+    }
+
+    functab = (comp == BD_METRIC_COMPONENT_PHIPHI) ?
+              eval_metric_phiphi : eval_metric_rhorho;
+
+    functab[diff_idx](bd, out, out_stride, nb_coords_rho, rho,
+                      nb_coords_z, z, psi, q);
+}
diff --git a/eval_metric_template.c b/eval_metric_template.c
new file mode 100644
index 0000000..8443785
--- /dev/null
+++ b/eval_metric_template.c
@@ -0,0 +1,68 @@
+/*
+ * Copyright 2016 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/>.
+ */
+
+/**
+ * @file
+ * the template for the loop over the grid points
+*/
+
+#define FUNC3(a, b, c) a ## _ ## b ## _ ## c
+#define FUNC2(a, b, c) FUNC3(a, b, c)
+#define FUNC(x) FUNC2(x, COMP, DIFF)
+
+#define DIFF_RHO (DIFF % 3)
+#define DIFF_Z   (DIFF / 3)
+
+static void FUNC(eval_metric)(const BDContext *bd, double *out, ptrdiff_t out_stride,
+                              int nb_coords_rho, const double *rho,
+                              int nb_coords_z,   const double *z,
+                              double *psi[9], double *q[9])
+{
+    BDPriv *s = bd->priv;
+
+    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 int idx = i * nb_coords_rho + j;
+            const double rrho = rho[j];
+
+            const double ppsi = psi[0][idx];
+            const double psi2 = SQR(ppsi);
+            const double psi3 = psi2 * ppsi;
+            const double psi4 = SQR(psi2);
+
+            const double base     = psi4 * exp(2 * q[0][idx]);
+            const double dqrho    = DIFF_RHO > 0 ? q[1][idx] + 2 * psi[1][idx] / ppsi : 0;
+            const double dqz      = DIFF_Z   > 0 ? q[3][idx] + 2 * psi[3][idx] / ppsi : 0;
+            const double dq2rho   = DIFF_RHO > 1 ? q[2][idx] + 2 * psi[2][idx] / ppsi - 2 * SQR(psi[1][idx] / ppsi) : 0;
+            const double dq2z     = DIFF_Z > 1 ? q[6][idx] + 2 * psi[6][idx] / ppsi - 2 * SQR(psi[3][idx] / ppsi) : 0;
+            const double dq2rho_z = DIFF_RHO && DIFF_Z ? q[4][idx] + 2 * psi[4][idx] / ppsi - 2 * psi[1][idx] * psi[3][idx] / psi2 : 0;
+
+            double val;
+
+            DO_EVAL;
+
+            *dst++ = val;
+        }
+    }
+}
+
+#undef FUNC3
+#undef FUNC2
+#undef FUNC
diff --git a/internal.h b/internal.h
index 78b3741..501d390 100644
--- a/internal.h
+++ b/internal.h
@@ -57,6 +57,14 @@ typedef struct BDPriv {
 
 int bdi_solve(BDContext *bd);
 
+void bdi_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],
+                     double *psi[9], double *q[9],
+                     double *out, ptrdiff_t out_stride);
+
 void bdi_log(const BDContext *bd, int level, const char *fmt, ...);
 void bdi_log_default_callback(const BDContext *bd, int level,
                               const char *fmt, va_list vl);