1 files changed, 77 insertions, 0 deletions

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;