Rewrite constraint evaluation code.

Make it much more efficient and easier to add other seed functions.

1 files changed, 77 insertions, 13 deletions

diff --git a/init.c b/init.c
index e90bbb8..d93d5a3 100644
--- a/init.c
+++ b/init.c
@@ -36,6 +36,7 @@
 #include "cpu.h"
 #include "log.h"
 #include "nlsolve.h"
+#include "td_constraints.h"
 #include "teukolsky_data.h"
 #include "threadpool.h"
 
@@ -241,16 +242,47 @@ static int teukolsky_init_check_options(TDContext *td)
     return 0;
 }
 
+static int constraint_eval_alloc(TDContext *td, double amplitude,
+                                 TDConstraintEvalContext **pce)
+{
+    TDPriv *priv = td->priv;
+    TDConstraintEvalContext *ce;
+    int ret;
+
+    ret = tdi_constraint_eval_alloc(&ce, TD_FAMILY_AE_TIME_ANTISYM);
+    if (ret < 0) {
+        tdi_log(&priv->logger, 0, "Error allocating the constraints evaluator\n");
+        return ret;
+    }
+
+    ce->logger       = priv->logger;
+    ce->amplitude    = amplitude;
+    ce->nb_coords[0] = td->nb_coeffs[0];
+    ce->nb_coords[1] = td->nb_coeffs[1];
+    ce->coords[0]    = priv->solver->colloc_grid[0][0];
+    ce->coords[1]    = priv->solver->colloc_grid[0][1];
+
+    ret = tdi_constraint_eval_init(ce);
+    if (ret < 0) {
+        tdi_constraint_eval_free(&ce);
+        tdi_log(&priv->logger, 0, "Error initializing the constraints evaluator\n");
+        return ret;
+    }
+
+    *pce = ce;
+    return 0;
+}
+
 static int teukolsky_constraint_eval(void *opaque, unsigned int eq_idx,
                                      const unsigned int *colloc_grid_order,
                                      const double * const *colloc_grid,
                                      const double * const (*vars)[PSSOLVE_DIFF_ORDER_NB],
                                      double *dst)
 {
+    TDConstraintEvalContext *ce = opaque;
     const double *amplitude = opaque;
 
-    return tdi_constraints_eval(eq_idx, *amplitude, colloc_grid_order, colloc_grid,
-                                vars, dst);
+    return tdi_constraint_eval(ce, eq_idx, vars, dst);
 }
 
 #define AMPLITUDE_DIRECT_SOLVE 1e-3
@@ -260,19 +292,26 @@ static int teukolsky_constraint_eval(void *opaque, unsigned int eq_idx,
 int td_solve(TDContext *td, double *coeffs_init[3])
 {
     TDPriv *s = td->priv;
+    TDConstraintEvalContext *ce;
+    double a0;
     int ret;
 
     ret = teukolsky_init_check_options(td);
     if (ret < 0)
         return ret;
 
-    if (td->amplitude > AMPLITUDE_CONVERGE) {
+    ret = constraint_eval_alloc(td, 0.0, &ce);
+    if (ret < 0)
+        return ret;
+    if (fabs(td->amplitude) >= ce->a_diverge) {
         tdi_log(&s->logger, 0,
-                "Amplitude A=%16.16g is above the branch point, no solutions "
+                "Amplitude A=%16.16g is above the point A_{max}=%g, no solutions "
                 "are known to exist there. Set solution_branch=1 to get to the "
                 "second branch where mass increases with decreasing amplitude\n",
-                td->amplitude);
+                td->amplitude, ce->a_converge);
     }
+    a0 = ce->a_converge;
+    tdi_constraint_eval_free(&ce);
 
     if (coeffs_init) {
         for (int i = 0; i < 3; i++) {
@@ -283,8 +322,13 @@ int td_solve(TDContext *td, double *coeffs_init[3])
 
     if (td->solution_branch == 0 || coeffs_init) {
         // direct solve with default (flat space) or user-provided initial guess
+        ret = constraint_eval_alloc(td, td->amplitude, &ce);
+        if (ret < 0)
+            return ret;
+
         ret = tdi_nlsolve_solve(s->solver, teukolsky_constraint_eval, NULL,
-                                &td->amplitude, s->coeffs, 0);
+                                ce, s->coeffs, 0);
+        tdi_constraint_eval_free(&ce);
         if (ret < 0) {
             tdi_log(&s->logger, 0, "tdi_nlsolve_solve() failed: %d", ret);
             return ret;
@@ -294,18 +338,27 @@ int td_solve(TDContext *td, double *coeffs_init[3])
         // execute two lower-branch solutions and extrapolate to get to the upper branch
         const int N = NB_EQUATIONS * td->nb_coeffs[0] * td->nb_coeffs[1];
         const double delta = 1e-5;
-        const double a0 = AMPLITUDE_CONVERGE;
-        const double a1 = AMPLITUDE_CONVERGE - delta;
+        const double a1 = a0 - delta;
 
         double cur_amplitude, new_amplitude, inverse_step;
         int dir;
 
+        ret = constraint_eval_alloc(td, a0, &ce);
+        if (ret < 0)
+            return ret;
+
         ret = tdi_nlsolve_solve(s->solver, teukolsky_constraint_eval, NULL,
-                                &a0, s->coeffs, 0);
+                                ce, s->coeffs, 0);
+        tdi_constraint_eval_free(&ce);
+        if (ret < 0)
+            return ret;
+
+        ret = constraint_eval_alloc(td, a1, &ce);
         if (ret < 0)
             return ret;
         ret = tdi_nlsolve_solve(s->solver, teukolsky_constraint_eval, NULL,
-                               &a1, s->coeffs_tmp, 0);
+                                ce, s->coeffs_tmp, 0);
+        tdi_constraint_eval_free(&ce);
         if (ret < 0)
             return ret;
 
@@ -313,8 +366,13 @@ int td_solve(TDContext *td, double *coeffs_init[3])
         cblas_daxpy(N, -1.0, s->coeffs_tmp, 1, s->coeffs,     1);
 
         // obtain solution for a1 in the upper branch
+        ret = constraint_eval_alloc(td, a1, &ce);
+        if (ret < 0)
+            return ret;
+
         ret = tdi_nlsolve_solve(s->solver, teukolsky_constraint_eval, NULL,
-                                &a1, s->coeffs, 0);
+                                ce, s->coeffs, 0);
+        tdi_constraint_eval_free(&ce);
         if (ret < 0) {
             tdi_log(&s->logger, 0, "Failed to get into the upper branch\n");
             return ret;
@@ -333,9 +391,15 @@ int td_solve(TDContext *td, double *coeffs_init[3])
             if (dir * (td->amplitude - new_amplitude) < 0)
                 new_amplitude = td->amplitude;
             tdi_log(&s->logger, 2, "Trying amplitude %g\n", new_amplitude);
+
+            ret = constraint_eval_alloc(td, new_amplitude, &ce);
+            if (ret < 0)
+                return ret;
+
             memcpy(s->coeffs_tmp, s->coeffs, sizeof(*s->coeffs) * N);
             ret = tdi_nlsolve_solve(s->solver, teukolsky_constraint_eval, NULL,
-                                    &new_amplitude, s->coeffs_tmp, 1);
+                                    ce, s->coeffs_tmp, 1);
+            tdi_constraint_eval_free(&ce);
             if (ret == -EDOM) {
                 inverse_step = 0.5 * inverse_step;
                 if (fabs(inverse_step) < 1e-2)
@@ -360,7 +424,7 @@ int td_solve(TDContext *td, double *coeffs_init[3])
 
     }
 finish:
-    tdi_solver_print_stats(s->solver);
+    tdi_nlsolve_print_stats(s->solver);
 
     return 0;
 }