Switch to polar coordinates in the solver.

1 files changed, 45 insertions, 1 deletions

diff --git a/basis.c b/basis.c
index bb3a8e9..cad7222 100644
--- a/basis.c
+++ b/basis.c
@@ -79,13 +79,29 @@ static double sb_even_eval_diff2(const BasisSetContext *s, double coord, int idx
     return sf * (idx + 1) * (2 * coord * cos((idx + 1) * val) - sf * (idx + 1) * sin((idx + 1) * val)) / SQR(SQR(sf) + SQR(coord));
 }
 
+static void sb_eval_multi(const BasisSetContext *s, double coord,
+                          double **out, int order_start, int order_end)
+{
+    const double sf = s->sf;
+    const double val = atan2(sf, coord);
+
+    for (int i = order_start; i < order_end; i++) {
+        const int idx = 2 * i;
+        const double sval = sin((idx + 1) * val);
+        const double cval = cos((idx + 1) * val);
+        out[0][i] = sval;
+        out[1][i] = - sf * (idx + 1) * cval / (SQR(sf) + SQR(coord));
+        out[2][i] = sf * (idx + 1) * (2 * coord * cval - sf * (idx + 1) * sval) / SQR(SQR(sf) + SQR(coord));
+    }
+}
+
 static double sb_even_colloc_point(const BasisSetContext *s, int order, int idx)
 {
     double t;
 
     idx = order - idx - 1;
 
-    t = (idx + 2) * M_PI / (2 * order + 2);
+    t = (idx + 2) * M_PI / (2 * order + 3);
     return s->sf / tan(t);
 }
 
@@ -98,6 +114,7 @@ static const BasisSet sb_even_basis = {
     .eval         = sb_even_eval,
     .eval_diff1   = sb_even_eval_diff1,
     .eval_diff2   = sb_even_eval_diff2,
+    .eval_multi   = sb_eval_multi,
     .colloc_point = sb_even_colloc_point,
 };
 
@@ -116,6 +133,18 @@ static double cos_even_eval_diff2(const BasisSetContext *s, double coord, int id
     return -4 * SQR(idx) * cos(2 * idx * coord);
 }
 
+static void cos_even_eval_multi(const BasisSetContext *s, double coord,
+                                double **out, int order_start, int order_end)
+{
+    for (int i = order_start; i < order_end; i++) {
+        const double sval = sin(2 * i * coord);
+        const double cval = cos(2 * i * coord);
+        out[0][i] = cval;
+        out[1][i] = -2 * i * sval;
+        out[2][i] = -4 * SQR(i) * cval;
+    }
+}
+
 static double cos_even_colloc_point(const BasisSetContext *s, int order, int idx)
 {
     return M_PI * idx / (2 * order);
@@ -125,6 +154,7 @@ static const BasisSet cos_even_basis = {
     .eval         = cos_even_eval,
     .eval_diff1   = cos_even_eval_diff1,
     .eval_diff2   = cos_even_eval_diff2,
+    .eval_multi   = cos_even_eval_multi,
     .colloc_point = cos_even_colloc_point,
 };
 
@@ -141,6 +171,20 @@ double bdi_basis_eval(BasisSetContext *s, enum BSEvalType type, double coord, in
     return eval(s, coord, order);
 }
 
+double bdi_basis_eval_multi(BasisSetContext *s, double coord, double **out,
+                            int order_start, int order_end)
+{
+    if (s->bs->eval_multi)
+        s->bs->eval_multi(s, coord, out, order_start, order_end);
+    else {
+        for (int i = order_start; i < order_end; i++) {
+            out[0][i] = s->bs->eval(s, coord, i);
+            out[1][i] = s->bs->eval_diff1(s, coord, i);
+            out[2][i] = s->bs->eval_diff2(s, coord, i);
+        }
+    }
+}
+
 double bdi_basis_colloc_point(BasisSetContext *s, int idx, int order)
 {
     return s->bs->colloc_point(s, order, idx);