TD_FAMILY_STRETCH WIP

controlled by env variables:
- TD_PARAM - param 'b'
- TD_A_MAX - max converging a value

3 files changed, 86 insertions, 2 deletions

diff --git a/init.c b/init.c
index 07a00d8..992feff 100644
--- a/init.c
+++ b/init.c
@@ -192,6 +192,7 @@ int td_solve(TDContext *td, double *coeffs_init[3])
     TDPriv *s = td->priv;
     NLSolveContext *nl;
     TDConstraintEvalContext *ce;
+    const char *env_a_max = getenv("TD_A_MAX");
     double a0;
     int ret;
 
@@ -206,14 +207,19 @@ int td_solve(TDContext *td, double *coeffs_init[3])
     ret = tdi_ce_alloc(td, td->nb_coeffs, nl->colloc_grid[0], 0.0, &ce);
     if (ret < 0)
         goto fail;
-    if (fabs(td->amplitude) >= ce->a_diverge) {
+    if (!env_a_max && fabs(td->amplitude) >= ce->a_diverge) {
         tdi_log(&s->logger, 0,
                 "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, ce->a_converge);
     }
-    a0 = SGN(td->amplitude) * ce->a_converge;
+
+    a0 = ce->a_converge;
+    if (env_a_max)
+        a0 = strtod(env_a_max, NULL);
+    a0 *= SGN(td->amplitude);
+
     tdi_constraint_eval_free(&ce);
 
     if (coeffs_init) {
diff --git a/td_constraints.c b/td_constraints.c
index ab21e79..f7201a6 100644
--- a/td_constraints.c
+++ b/td_constraints.c
@@ -17,6 +17,7 @@
 
 #include <errno.h>
 #include <math.h>
+#include <stdlib.h>
 
 #include "common.h"
 #include "pssolve.h"
@@ -25,6 +26,7 @@
 struct TDConstraintEvalPriv {
     double *k_rtheta;
     double *dk_rtheta[2];
+    double param;
 };
 
 typedef void (*ConstraintEvalCB)(TDConstraintEvalContext *ctx,
@@ -179,6 +181,76 @@ static const ConstraintEvalCB constraint_funcs_confflat[TD_CONSTRAINT_EQ_NB] = {
     [TD_CONSTRAINT_EQ_MOM_1] = constraint_eval_mom_t_confflat,
 };
 
+static double k_rtheta_eval_stretch(TDConstraintEvalContext *ctx, double r, double theta)
+{
+    const double a = ctx->amplitude;
+    const double b = ctx->priv->param;
+
+    const double st = sin(theta);
+    const double ct = cos(theta);
+
+    const double  r2 = r * r;
+    const double  b2 = b * b;
+    const double st2 = st * st;
+    const double ct2 = ct * ct;
+
+    return -a * exp(-b2 - r2) * r * st * (-2 * ct * cosh(2 * b * r * ct) * (-1.0 + r2 + b2 * r2 * st2) +
+                                          b * r * (2 * ct2 + (-3. + 2. * r2) * st2) * sinh(2 * b * r * ct));
+}
+
+static double k_rtheta_eval_dr_stretch(TDConstraintEvalContext *ctx, double r, double theta)
+{
+    const double a = ctx->amplitude;
+    const double b = ctx->priv->param;
+
+    const double st  = sin(theta);
+    const double ct  = cos(theta);
+    const double c2t = cos(2. * theta);
+
+    const double  r2 = r * r;
+    const double  r4 = r2 * r2;
+
+    const double  b2 = b * b;
+    const double st2 = st * st;
+    const double ct2 = ct * ct;
+
+    return 2. * a * exp(-b2 - r2) * st * (ct * (-1. + (5. + 2 * b2) * r2 - 2. * (1. + b2) * r4 + 2. * b2 * r2 * (-2. + r2) * c2t) * cosh(2 * b * r * ct) +
+                                          b * r * ((3. - 7. * r2 + 2. * r4) * st2 + 2. * ct2 * (-2. + 2. * r2 + b2 * r2 * st2)) * sinh(2 * b * r * ct));
+}
+
+static double k_rtheta_eval_dt_stretch(TDConstraintEvalContext *ctx, double r, double theta)
+{
+    const double a = ctx->amplitude;
+    const double b = ctx->priv->param;
+
+    const double st = sin(theta);
+    const double ct = cos(theta);
+
+    const double  r2 = r * r;
+    const double  r4 = r2 * r2;
+
+    const double  b2 = b * b;
+    const double st2 = st * st;
+    const double ct2 = ct * ct;
+
+    return -a * exp(-b2 - r2) * r * (-2. * cosh(2 * b * r * ct) * (ct2 * (-1. + r2 + 5. * b2 * r2 * st2) + st2 * (1. - r2 + 2 * b2 * r2 * (-2. + r2) * st2)) +
+                                     b * r * ct * (2. * ct2 + st2 * (-17.0 + 10.0 * r2 + 4. * b2 * r2 * st2)) * sinh(2 * b * r * ct));
+}
+
+static const TDFamilyDef stretch = {
+    .eval_krt    = k_rtheta_eval_stretch,
+    .eval_krt_dr = k_rtheta_eval_dr_stretch,
+    .eval_krt_dt = k_rtheta_eval_dt_stretch,
+    // b = 1.0
+    //.a_converge  = 1.996975898742676,
+    //.a_diverge   = 1.996976852416992,
+    // b=0.5
+    //.a_converge = 1.711344718933105,
+    .a_converge = 1.711344,
+    .a_diverge  = 1.711345672607422,
+    .constraint_eval = constraint_funcs_confflat,
+};
+
 static double k_rtheta_eval_time_antisym_cubic(TDConstraintEvalContext *ctx,
                                                double r, double theta)
 {
@@ -247,6 +319,7 @@ static const TDFamilyDef time_antisym_linear = {
 static const TDFamilyDef *td_families[] = {
     [TD_FAMILY_TIME_ANTISYM_CUBIC] = &time_antisym_cubic,
     [TD_FAMILY_TIME_ANTISYM_LINEAR] = &time_antisym_linear,
+    [TD_FAMILY_STRETCH]             = &stretch,
 };
 
 double tdi_constraint_eval_k_rtheta(TDConstraintEvalContext *ctx, double r, double theta)
@@ -272,6 +345,7 @@ int tdi_constraint_eval_init(TDConstraintEvalContext *ctx)
     TDConstraintEvalPriv *priv = ctx->priv;
     const TDFamilyDef *fd = td_families[ctx->family];
     int ret;
+    const char *param = getenv("TD_PARAM");
 
     ret = posix_memalign((void**)&priv->k_rtheta, 32,
                          sizeof(*priv->k_rtheta) * ctx->nb_coords[0] * ctx->nb_coords[1]);
@@ -285,6 +359,9 @@ int tdi_constraint_eval_init(TDConstraintEvalContext *ctx)
             return -ENOMEM;
     }
 
+    if (param)
+        priv->param = strtod(param, NULL);
+
     ctx->a_converge = fd->a_converge;
     ctx->a_diverge  = fd->a_diverge;
 
diff --git a/teukolsky_data.h b/teukolsky_data.h
index 3569e61..f4d8244 100644
--- a/teukolsky_data.h
+++ b/teukolsky_data.h
@@ -59,6 +59,7 @@ enum TDFamily {
      *  θ       L       \    L     /
      */
     TD_FAMILY_TIME_ANTISYM_LINEAR,
+    TD_FAMILY_STRETCH,
 };
 
 typedef struct TDContext {