td_constraints: fix constraint evaluation at the axis

1 files changed, 24 insertions, 3 deletions

diff --git a/td_constraints.c b/td_constraints.c
index 1c63954..aba87ab 100644
--- a/td_constraints.c
+++ b/td_constraints.c
@@ -52,6 +52,10 @@ constraint_eval_ham_confflat(TDConstraintEvalContext *ctx,
         const double theta = ctx->coords[1][idx1];
         const double tt = tan(theta);
 
+        const double tt_normalized = tt / M_PI;
+        const int    is_axis       = ABS(tt_normalized - round(tt_normalized)) < EPS;
+        const double dtt_t         = ((int)tt_normalized) & 1 ? -1.0 : 1.0;
+
         for (int idx0 = 0; idx0 < ctx->nb_coords[0]; idx0++) {
             const double r = ctx->coords[0][idx0];
 
@@ -74,7 +78,10 @@ constraint_eval_ham_confflat(TDConstraintEvalContext *ctx,
             const double r2   = SQR(r);
             const double psi2 = SQR(psi);
             const double psi4 = SQR(psi2);
-            const double Rscal = -8.0 * (d2psi_rr + 2.0 * dpsi_r / r + d2psi_tt / r2 + dpsi_t / (r2 * tt)) / (psi4 * psi);
+
+            // term that is singular on the axis, use l'Hospital rule
+            const double dpsi_term = is_axis ? d2psi_tt / dtt_t : dpsi_t / tt;
+            const double Rscal = -8.0 * (d2psi_rr + 2.0 * dpsi_r / r + d2psi_tt / r2 + dpsi_term / r2) / (psi4 * psi);
 
             double K2 = 0.0;
             for (int i = 0; i < 3; i++)
@@ -97,6 +104,10 @@ constraint_eval_mom_r_confflat(TDConstraintEvalContext *ctx,
         const double theta = ctx->coords[1][idx1];
         const double tt = tan(theta);
 
+        const double tt_normalized = tt / M_PI;
+        const int    is_axis       = ABS(tt_normalized - round(tt_normalized)) < EPS;
+        const double dtt_t         = ((int)tt_normalized) & 1 ? -1.0 : 1.0;
+
         for (int idx0 = 0; idx0 < ctx->nb_coords[0]; idx0++) {
             const double r = ctx->coords[0][idx0];
 
@@ -113,8 +124,11 @@ constraint_eval_mom_r_confflat(TDConstraintEvalContext *ctx,
 
             const double r2   = SQR(r);
 
+            // term that is singular on the axis, use l'Hospital
+            const double Krt_singular = is_axis ? dK_rt_t / dtt_t : K_rt / tt;
+
             out[idx] = 6.0 * (K_rr * dpsi_r + K_rt * dpsi_t / r2) / psi +
-                       dK_rr_r + 3.0 * K_rr / r + K_rt / (r2 * tt) + dK_rt_t / r2;
+                       dK_rr_r + 3.0 * K_rr / r + Krt_singular / r2 + dK_rt_t / r2;
         }
     }
 }
@@ -130,6 +144,10 @@ constraint_eval_mom_t_confflat(TDConstraintEvalContext *ctx,
         const double theta = ctx->coords[1][idx1];
         const double tt = tan(theta);
 
+        const double tt_normalized = tt / M_PI;
+        const int    is_axis       = ABS(tt_normalized - round(tt_normalized)) < EPS;
+        const double dtt_t         = ((int)tt_normalized) & 1 ? -1.0 : 1.0;
+
         for (int idx0 = 0; idx0 < ctx->nb_coords[0]; idx0++) {
             const double r = ctx->coords[0][idx0];
 
@@ -146,7 +164,10 @@ constraint_eval_mom_t_confflat(TDConstraintEvalContext *ctx,
             const double dpsi_r   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_10][idx];
             const double dpsi_t   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_01][idx];
 
-            out[idx] = -(K_rr + 2.0 * K_pp) / tt - 6.0 * (K_rr + K_pp) * dpsi_t / psi +
+            // term that is singular on the axis, use l'Hospital
+            const double K_singular = is_axis ? (dK_rr_t + 2.0 * dK_pp_t) / dtt_t : (K_rr + 2.0 * K_pp) / tt;
+
+            out[idx] = -K_singular - 6.0 * (K_rr + K_pp) * dpsi_t / psi +
                        6.0 * K_rt * dpsi_r / psi - dK_pp_t - dK_rr_t + dK_rt_r + 2.0 * K_rt / r;
         }
     }