Add a function for evaluating the metric.

2 files changed, 297 insertions, 14 deletions

diff --git a/brill_data.c b/brill_data.c
index 159d7c0..72a44a3 100644
--- a/brill_data.c
+++ b/brill_data.c
@@ -56,11 +56,18 @@ typedef struct BasisSet {
     long double (*colloc_point)(int order, int idx, long double sf);
 } BasisSet;
 
+typedef struct QFunc {
+    double         (*q)(struct BDContext *bd, double rho, double z);
+    double    (*dq_rho)(struct BDContext *bd, double rho, double z);
+    double      (*dq_z)(struct BDContext *bd, double rho, double z);
+    double   (*d2q_rho)(struct BDContext *bd, double rho, double z);
+    double     (*d2q_z)(struct BDContext *bd, double rho, double z);
+    double (*d2q_rho_z)(struct BDContext *bd, double rho, double z);
+} QFunc;
+
 typedef struct BDPriv {
     const BasisSet *basis;
-
-    double   (*q_func)(struct BDContext *bd, double rho, double z);
-    double (*d2q_func)(struct BDContext *bd, double rho, double z);
+    const QFunc    *q_func;
 
     int nb_colloc_points;
     int nb_colloc_points_rho;
@@ -152,7 +159,7 @@ static int brill_construct_matrix(BDContext *bd, gsl_matrix *mat,
 
         for (idx_grid_rho = 0; idx_grid_rho < s->nb_colloc_points_rho; idx_grid_rho++) {
             double x_physical = s->basis->colloc_point(s->colloc_grid_order_rho, idx_grid_rho, sfr);
-            double d2q        = s->d2q_func(bd, x_physical, z_physical);
+            double d2q        = s->q_func->d2q_rho(bd, x_physical, z_physical) + s->q_func->d2q_z(bd, x_physical, z_physical);
             int idx_grid         = idx_grid_z * s->nb_colloc_points_z + idx_grid_rho;
 
             for (idx_coeff_z = 0; idx_coeff_z < bd->nb_coeffs_z; idx_coeff_z++)
@@ -335,21 +342,132 @@ static double q_gundlach(BDContext *bd, double rho, double z)
     return bd->amplitude * SQR(rho) * exp(- (SQR(rho) + SQR(z)));
 }
 
-static double d2q_gundlach(BDContext *bd, double rho, double z)
+static double dq_rho_gundlach(BDContext *bd, double rho, double z)
+{
+    return bd->amplitude * 2 * rho * exp(-SQR(rho) - SQR(z)) * (1 - SQR(rho));
+}
+
+static double d2q_rho_gundlach(BDContext *bd, double rho, double z)
+{
+    double rho2 = SQR(rho);
+    return bd->amplitude * 2 * exp(-rho2 - SQR(z)) * ((1 - rho2) * (1 - 2 * rho2) - 2 * rho2);
+}
+
+static double d2q_rho_z_gundlach(BDContext *bd, double rho, double z)
+{
+    return -bd->amplitude * 4 * z * rho * exp(-SQR(rho) - SQR(z)) * (1 - SQR(rho));
+}
+
+static double dq_z_gundlach(BDContext *bd, double rho, double z)
 {
-    double r2 = SQR(rho);
-    double z2 = SQR(z);
-    double e  = exp(-r2 - z2);
+    return - bd->amplitude * 2 * z * SQR(rho) * exp(-SQR(rho) - SQR(z));
+}
 
-    return 2 * bd->amplitude * e * (1.0 + 2 * r2 * (r2 + z2 - 3));
+static double d2q_z_gundlach(BDContext *bd, double rho, double z)
+{
+    return bd->amplitude * 2 * SQR(rho) * exp(-SQR(rho) - SQR(z)) * (2 * SQR(z) - 1);
 }
 
+static const QFunc q_func_gundlach = {
+    .q         = q_gundlach,
+    .dq_rho    = dq_rho_gundlach,
+    .dq_z      = dq_z_gundlach,
+    .d2q_rho   = d2q_rho_gundlach,
+    .d2q_z     = d2q_z_gundlach,
+    .d2q_rho_z = d2q_rho_z_gundlach,
+};
+
 static double q_eppley(BDContext *bd, double rho, double z)
 {
     double r2 = SQR(rho) + SQR(z);
     return bd->amplitude * SQR(rho) / (1.0 + pow(r2, bd->eppley_n / 2.0));
 }
 
+static double dq_rho_eppley(BDContext *bd, double rho, double z)
+{
+    double A = bd->amplitude;
+    double n = bd->eppley_n;
+    double rho2 = SQR(rho);
+    double z2   = SQR(z);
+
+    double r2 = rho2 + z2;
+    double r  = sqrt(r2);
+    double rnm2 = pow(r, n - 2);
+
+    return A * rho * (2 * (1 + rnm2 * r2) - n * rho2 * rnm2) / SQR(1 + rnm2 * r2);
+}
+
+static double dq_z_eppley(BDContext *bd, double rho, double z)
+{
+    double A = bd->amplitude;
+    double n = bd->eppley_n;
+    double rho2 = SQR(rho);
+    double z2   = SQR(z);
+
+    double r2 = rho2 + z2;
+    double r  = sqrt(r2);
+    double rnm2 = pow(r, n - 2);
+
+    return - A * n * rho2 * z * rnm2 / SQR(1 + rnm2 * r2);
+}
+
+static double d2q_rho_eppley(BDContext *bd, double rho, double z)
+{
+    double A = bd->amplitude;
+    double n = bd->eppley_n;
+    double rho2 = SQR(rho);
+    double z2   = SQR(z);
+
+    double r2 = rho2 + z2;
+    double r  = sqrt(r2);
+    double rnm4 = pow(r, n - 4);
+    double rn   = rnm4 * SQR(r2);
+    double rnp1 = 1 + rn;
+
+    return A * (SQR(n) * SQR(rho2) * rnm4 * (rn - 1) - n * rho2 * rnm4 * (3 * rho2 + 5 * z2) * rnp1 + 2 * SQR(rnp1)) / (rnp1 * SQR(rnp1));
+}
+
+static double d2q_z_eppley(BDContext *bd, double rho, double z)
+{
+    double A = bd->amplitude;
+    double n = bd->eppley_n;
+    double rho2 = SQR(rho);
+    double z2   = SQR(z);
+
+    double r2 = rho2 + z2;
+    double r  = sqrt(r2);
+    double rnm4 = pow(r, n - 4);
+    double rn   = rnm4 * SQR(r2);
+    double rnp1 = 1 + rn;
+
+    return A * n * rho2 * rnm4 * (z2 - rho2 - n * z2 + rn * ((1 + n) * z2 - rho2)) / (rnp1 * SQR(rnp1));
+}
+
+static double d2q_rho_z_eppley(BDContext *bd, double rho, double z)
+{
+    double A = bd->amplitude;
+    double n = bd->eppley_n;
+    double rho2 = SQR(rho);
+    double z2   = SQR(z);
+
+    double r2 = rho2 + z2;
+    double r  = sqrt(r2);
+    double rnm4 = pow(r, n - 4);
+    double rn   = rnm4 * SQR(r2);
+    double rnp1 = 1 + rn;
+
+    return A * n * rho * z * rnm4 * (rn * (n * rho2 - 2 * z2) - 2 * z2 - n * rho2) / (rnp1 * SQR(rnp1));
+}
+
+static const QFunc q_func_eppley = {
+    .q         = q_eppley,
+    .dq_rho    = dq_rho_eppley,
+    .dq_z      = dq_z_eppley,
+    .d2q_rho   = d2q_rho_eppley,
+    .d2q_z     = d2q_z_eppley,
+    .d2q_rho_z = d2q_rho_z_eppley,
+};
+
 static double d2q_eppley(BDContext *bd, double rho, double z)
 {
     double rho2 = SQR(rho);
@@ -370,17 +488,15 @@ static int brill_init_check_options(BDContext *bd)
 
     switch (bd->q_func_type) {
     case BD_Q_FUNC_GUNDLACH:
-        s->q_func   = q_gundlach;
-        s->d2q_func = d2q_gundlach;
+        s->q_func   = &q_func_gundlach;
         break;
     case BD_Q_FUNC_EPPLEY:
-        s->q_func   = q_eppley;
-        s->d2q_func = d2q_eppley;
-
         if (bd->eppley_n < 4) {
             fprintf(stderr, "Invalid n: %d < 4\n", bd->eppley_n);
             return -EINVAL;
         }
+
+        s->q_func   = &q_func_eppley;
         break;
     default:
         fprintf(stderr, "Unknown q function type: %d\n", bd->q_func_type);
@@ -546,6 +662,148 @@ int bd_eval_psi(BDContext *bd, const double *rho, int nb_coords_rho,
 
 }
 
+int bd_eval_metric(BDContext *bd, const double *rho, int nb_coords_rho,
+                   const double *z, int nb_coords_z,
+                   const unsigned int comp[2], const unsigned int diff_order[2],
+                   double *out)
+{
+    BDPriv *s = bd->priv;
+    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;
+    int ret = 0;
+
+    /* check the parameters for validity */
+    if (comp[0] > 2 || comp[1] > 2) {
+        fprintf(stderr, "Invalid component indices: %d%d\n", comp[0], comp[1]);
+        return -EINVAL;
+    }
+
+    if (comp[0] != comp[1]) {
+        memset(out, 0, nb_coords * sizeof(*out));
+        return 0;
+    }
+
+    if (diff_order[0] + diff_order[1] > 2) {
+        fprintf(stderr, "At most second order derivatives are supported.\n");
+        return -ENOSYS;
+    }
+
+    /* evaluate the conformal factor and its derivatives if necessary */
+#define EVAL_PSI(arr, order)                                    \
+do {                                                            \
+    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);                              \
+    if (ret < 0)                                                \
+        goto fail;                                              \
+} while (0)
+
+    EVAL_PSI(psi,             ((unsigned int[2]){ 0, 0 }));
+    if (diff_order[0])
+        EVAL_PSI(dpsi_rho,    ((unsigned int[2]){ 1, 0 }));
+    if (diff_order[0] > 1)
+        EVAL_PSI(d2psi_rho,   ((unsigned int[2]){ 2, 0 }));
+    if (diff_order[1])
+        EVAL_PSI(dpsi_z,      ((unsigned int[2]){ 0, 1 }));
+    if (diff_order[1] > 1)
+        EVAL_PSI(d2psi_z,     ((unsigned int[2]){ 0, 2 }));
+    if (diff_order[0] && diff_order[1])
+        EVAL_PSI(d2psi_rho_z, ((unsigned int[2]){ 1, 1 }));
+
+/* 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];                                     \
+                                                                    \
+        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;                                                \
+                                                                    \
+            out[i * nb_coords_rho + j] = val;                       \
+        }                                                           \
+    }                                                               \
+} while (0)
+
+#define GRID(arr) (arr[i * nb_coords_rho + j])
+
+    if (comp[0] == 2) {
+        // γ_φφ
+        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 * s->q_func->q(bd, rrho, zz));                                                                          \
+    double drho, dz, d2rho, d2z, d2rho_z;                                                                                                    \
+                                                                                                                                             \
+    if (eval_drho)   drho    = s->q_func->dq_rho(bd, rrho, zz)    + 2 * GRID(dpsi_rho)    / ppsi;                                            \
+    if (eval_dz)     dz      = s->q_func->dq_z(bd, rrho, zz)      + 2 * GRID(dpsi_z)      / ppsi;                                            \
+    if (eval_d2rho)  d2rho   = s->q_func->d2q_rho(bd, rrho, zz)   + 2 * GRID(d2psi_rho)   / ppsi - 2 * SQR(GRID(dpsi_rho) / ppsi);           \
+    if (eval_d2z)    d2z     = s->q_func->d2q_z(bd, rrho, zz)     + 2 * GRID(d2psi_z)     / ppsi - 2 * SQR(GRID(dpsi_z)   / ppsi);           \
+    if (eval_d2rhoz) d2rho_z = s->q_func->d2q_rho_z(bd, rrho, zz) + 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;        // ∂_ρ ∂_ρ
+        }
+    }
+
+fail:
+    free(psi);
+    free(dpsi_rho);
+    free(dpsi_z);
+    free(d2psi_rho);
+    free(d2psi_rho_z);
+    free(d2psi_z);
+
+    return ret;
+}
+
 #if 0
 void brill_data(CCTK_ARGUMENTS)
 {
diff --git a/brill_data.h b/brill_data.h
index 85e1aac..a3f897b 100644
--- a/brill_data.h
+++ b/brill_data.h
@@ -161,3 +161,28 @@ int bd_eval_psi(BDContext *bd, const double *rho, int nb_coords_rho,
                 const double *z, int nb_coords_z,
                 const unsigned int diff_order[2],
                 double *psi);
+
+/**
+ * Evaluate the 3-metric γ_ij at the specified rectangular grid (in cylindrical
+ * coordinates { ρ, z, φ }).
+ *
+ * @param bd the solver context
+ * @param rho the array of ρ coordinates.
+ * @param nb_coords_rho the number of elements in rho.
+ * @param z   the array of z coordinates.
+ * @param nb_coords_z the number of elements in z.
+ * @param comp the component of the metric to evaluate.
+ * @param diff_order the order of the derivatives of the metric to evaluate. The first element
+ *                   specifies the derivative wrt ρ, the second wrt z. I.e. diff_order = { 0, 0 }
+ *                   evaluates the metric itself, diff_order = { 0, 1 } evaluates ∂γ/∂z etc.
+ * @param out the array into which the values of the metric will be written. The metric
+ *            is evaluated on the grid formed by the outer product of the rho
+ *            and z vectors. I.e. out[j * nb_coords_rho + i] = γ_comp[0]comp[1](rho[i], z[j]).
+ *            The length of out must be nb_coords_rho * nb_coords_z.
+ *
+ * @return >= 0 on success, a negative error code on failure.
+ */
+int bd_eval_metric(BDContext *bd, const double *rho, int nb_coords_rho,
+                   const double *z, int nb_coords_z,
+                   const unsigned int comp[2], const unsigned int diff_order[2],
+                   double *out);