From 1f443d3fdb93e71a70244b7521ae0f4c40601351 Mon Sep 17 00:00:00 2001
From: Anton Khirnov <anton@khirnov.net>
Date: Fri, 27 Aug 2021 13:35:52 +0200
Subject: Add public API for evaluating the residual at an arbitrary location.

---
 init.c | 87 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 82 insertions(+), 5 deletions(-)

(limited to 'init.c')

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;
+}
-- 
cgit v1.2.3