1 files changed, 64 insertions, 0 deletions
diff --git a/lib/lambert_w.c b/lib/lambert_w.c
new file mode 100644
index 0000000..da59199
--- /dev/null
+++ b/lib/lambert_w.c
@@ -0,0 +1,64 @@
+// Lambert W
+//
+// lambert_w
+//
+// Lambert W function
+//
+// W(x) such that W*e^W = x
+//
+// Inputs:    1
+// Outputs:   1
+// Constants: 0
+
+
+#include <math.h>
+#include "ts.h"
+
+
+ts_parameter_t const ts_lambert_w_p = { .n_inputs = 1, .n_outputs = 1, .n_constants = 0, .params_size = 0 };
+
+uint_t ts_lambert_w_offset( ts_constant_t const *c )
+{
+	( void )c;
+	return ( 0 );
+}
+
+int ts_lambert_w(
+	fp_t * const *y, uint_t y_n,
+	fp_t const * const *x, uint_t x_n,
+	uint_t n, ts_constant_t const *c
+)
+{
+	TS_PROLOGUE( ts_lambert_w_p );
+	(void)c;
+    for (uint_t i = 0; i < n; ++i)
+    {
+        fp_t xv = x[0][i];
+        if (xv < -1.0/exp(1.0))
+        {
+            y[0][i] = NAN;
+            continue;
+        }
+        if (xv == 0.0)
+        {
+            y[0][i] = 0.0;
+            continue;
+        }
+        /* Halley iteration for W(x) */
+        fp_t w = xv < 1.0 ? log(xv + 1.0) : log(xv);
+        for (uint_t j = 0; j < 10; ++j)
+        {
+            fp_t ew = exp(w);
+            fp_t f = w * ew - xv;
+            fp_t f1 = ew * (w + 1.0);
+            if (f1 == 0) break;
+            fp_t f2 = f1 * (w + 2.0);
+            w = w - f / f1 * (1.0 + f * f2 / (2.0 * f1 * f1));
+        }
+        y[0][i] = w;
+    }
+	return 0;
+}
+
+
+