Add public API for evaluating the residual at an arbitrary location.

5 files changed, 114 insertions, 6 deletions

diff --git a/common.h b/common.h
index 16dd7ef..55359ae 100644
--- a/common.h
+++ b/common.h
@@ -18,6 +18,7 @@
 #ifndef TEUKOLSKY_DATA_COMMON_H
 #define TEUKOLSKY_DATA_COMMON_H
 
+#define ABS(x) (((x) >= 0) ? (x) : -(x))
 #define SQR(x) ((x) * (x))
 #define SGN(x) ((x) >= 0.0 ? 1.0 : -1.0)
 #define MAX(x, y) ((x) > (y) ? (x) : (y))
diff --git a/init.c b/init.c
index c5c633c..dac0731 100644
--- a/init.c
+++ b/init.c
@@ -796,7 +796,7 @@ static int eval_var_kernel(void *arg, unsigned int job_idx, unsigned int thread_
 
 static int eval_var(const TDContext *td, unsigned int var_idx,
                     size_t nb_coords, const double *r, const double *theta,
-                    const unsigned int diff_order[2],
+                    const unsigned int diff_order[2], double add,
                     double *out)
 {
     TDPriv *s = td->priv;
@@ -830,7 +830,7 @@ static int eval_var(const TDContext *td, unsigned int var_idx,
     thread_data.basis[0]      = s->basis[var_idx][0];
     thread_data.basis[1]      = s->basis[var_idx][1];
     thread_data.coeffs        = td->coeffs[var_idx];
-    thread_data.init_val      = (var_idx == 0 && diff_order[0] == 0 && diff_order[1] == 0) ? 1.0 : 0.0;
+    thread_data.init_val      = add;
 
     tp_execute(s->tp, nb_coords, eval_var_kernel, &thread_data);
 
@@ -846,7 +846,8 @@ int td_eval_psi(const TDContext *td,
                 const unsigned int diff_order[2],
                 double *out)
 {
-    return eval_var(td, 0, nb_coords, r, theta, diff_order, out);
+    const double add = (diff_order[0] == 0 && diff_order[1] == 0) ? 1.0 : 0.0;
+    return eval_var(td, 0, nb_coords, r, theta, diff_order, add, out);
 }
 
 int td_eval_krr(const TDContext *td,
@@ -854,14 +855,14 @@ int td_eval_krr(const TDContext *td,
                 const unsigned int diff_order[2],
                 double *out)
 {
-    return eval_var(td, 1, nb_coords, r, theta, diff_order, out);
+    return eval_var(td, 1, nb_coords, r, theta, diff_order, 0.0, out);
 }
 int td_eval_kpp(const TDContext *td,
                 size_t nb_coords, const double *r, const double *theta,
                 const unsigned int diff_order[2],
                 double *out)
 {
-    return eval_var(td, 2, nb_coords, r, theta, diff_order, out);
+    return eval_var(td, 2, nb_coords, r, theta, diff_order, 0.0, out);
 }
 
 int td_eval_krt(const TDContext *td,
@@ -955,3 +956,79 @@ fail:
 
     return 0;
 }
+
+int td_eval_residual(const TDContext *td,
+                     size_t nb_coords_r,     const double *r,
+                     size_t nb_coords_theta, const double *theta,
+                     double *out[3])
+{
+    const unsigned int nb_coords[2] = { nb_coords_r, nb_coords_theta };
+    const double         *coords[2] = { r, theta};
+    const size_t          grid_size = nb_coords_r * nb_coords_theta;
+
+    TDConstraintEvalContext *ce = NULL;
+    double *r_2d = NULL, *theta_2d = NULL;
+    double *vars[TD_CONSTRAINT_VAR_NB][PSSOLVE_DIFF_ORDER_NB] = {{ NULL }};
+    int ret;
+
+    ret  = posix_memalign((void**)&r_2d,     32, grid_size * sizeof(*r_2d));
+    ret |= posix_memalign((void**)&theta_2d, 32, grid_size * sizeof(*theta_2d));
+    if (ret) {
+        ret = -ENOMEM;
+        goto fail;
+    }
+
+    for (size_t j = 0; j < nb_coords_theta; j++) {
+        for (size_t i = 0; i < nb_coords_r; i++) {
+            const size_t idx = j * nb_coords_r + i;
+            r_2d[idx]     = r[i];
+            theta_2d[idx] = theta[j];
+        }
+    }
+
+    for (int var = 0; var < TD_CONSTRAINT_VAR_NB; var++) {
+        for (int diff_order = 0; diff_order < PSSOLVE_DIFF_ORDER_NB; diff_order++) {
+            const unsigned int diff_orders[2] = {
+                tdi_pssolve_diff_order(diff_order, 0),
+                tdi_pssolve_diff_order(diff_order, 1),
+            };
+
+            ret  = posix_memalign((void**)&vars[var][diff_order],
+                                  32, grid_size * sizeof(*vars[var][diff_order]));
+            if (ret) {
+                ret = -ENOMEM;
+                goto fail;
+            }
+
+            ret = eval_var(td, var, grid_size, r_2d, theta_2d, diff_orders,
+                           0.0, vars[var][diff_order]);
+            if (ret)
+                goto fail;
+        }
+    }
+
+    ret = constraint_eval_alloc(td, nb_coords, coords,
+                                td->amplitude, &ce);
+    if (ret < 0)
+        goto fail;
+
+    ret  = tdi_constraint_eval(ce, TD_CONSTRAINT_EQ_HAM,   vars, out[0]);
+    ret |= tdi_constraint_eval(ce, TD_CONSTRAINT_EQ_MOM_0, vars, out[1]);
+    ret |= tdi_constraint_eval(ce, TD_CONSTRAINT_EQ_MOM_1, vars, out[2]);
+    if (ret)
+        goto fail;
+
+    ret = 0;
+
+fail:
+    free(r_2d);
+    free(theta_2d);
+    tdi_constraint_eval_free(&ce);
+    for (int diff_order = 0; diff_order < PSSOLVE_DIFF_ORDER_NB; diff_order++) {
+        free(vars[0][diff_order]);
+        free(vars[1][diff_order]);
+        free(vars[2][diff_order]);
+    }
+
+    return ret;
+}
diff --git a/libteukolskydata.v b/libteukolskydata.v
index 7d01e87..168f2eb 100644
--- a/libteukolskydata.v
+++ b/libteukolskydata.v
@@ -1,4 +1,4 @@
-LIBTEUKOLSKYDATA_3 {
+LIBTEUKOLSKYDATA_4 {
     global: td_*;
     local: *;
 };
diff --git a/teukolsky_data.h b/teukolsky_data.h
index 786554e..861613d 100644
--- a/teukolsky_data.h
+++ b/teukolsky_data.h
@@ -207,4 +207,13 @@ int td_eval_lapse(const TDContext *td,
                   const unsigned int diff_order[2],
                   double *out);
 
+/**
+ * Evaluate the residual of the three constraint equations at the specified
+ * coordinates.
+ */
+int td_eval_residual(const TDContext *td,
+                     size_t nb_coords_r,     const double *r,
+                     size_t nb_coords_theta, const double *theta,
+                     double *out[3]);
+
 #endif /* TEUKOLSKY_DATA_H */
diff --git a/teukolsky_data.py b/teukolsky_data.py
index 00ec5c4..a9bd490 100644
--- a/teukolsky_data.py
+++ b/teukolsky_data.py
@@ -170,6 +170,27 @@ class TeukolskyData(object):
 
         return metric, curv
 
+    def eval_residual(self, r, theta):
+        if r.ndim != 1 or theta.ndim != 1:
+            raise TypeError('r and theta must both be 1-dimensional NumPy arrays')
+
+        out = np.empty((3, r.shape[0] * theta.shape[0]))
+
+        c_out    = (ctypes.POINTER(ctypes.c_double) * 3)()
+        for i in range(3):
+            c_out[i] = ctypes.cast(np.ctypeslib.as_ctypes(out[i]), ctypes.POINTER(ctypes.c_double))
+
+        c_r     = np.ctypeslib.as_ctypes(r)
+        c_theta = np.ctypeslib.as_ctypes(theta)
+
+        ret = self._libtd.td_eval_residual(self._tdctx, r.shape[0], c_r,
+                                           theta.shape[0], c_theta, c_out)
+        if ret < 0:
+            raise RuntimeError('Error evaluating the variable: %d' % ret)
+
+        out.shape = (3, theta.shape[0], r.shape[0])
+        return out
+
     @property
     def amplitude(self):
         return self._tdctx.contents.__getattribute__('amplitude')