From 45c17fb90975f8ec7d0743a52bb0ea78d371424b Mon Sep 17 00:00:00 2001
From: Vanya Belyaev <ivan.belyaev@cern.ch>
Date: Sun, 27 Oct 2024 15:28:51 +0100
Subject: [PATCH]  Add Lambert W function

---
 source/include/Ostap/MoreMath.h       | 13 ++++++++
 source/include/Ostap/ValueWithError.h |  5 ++++
 source/src/MoreMath.cpp               | 43 +++++++++++++++++++++++++++
 source/src/ValueWithError.cpp         | 18 +++++++++++
 4 files changed, 79 insertions(+)

diff --git a/source/include/Ostap/MoreMath.h b/source/include/Ostap/MoreMath.h
index bd86afa5..153aec42 100644
--- a/source/include/Ostap/MoreMath.h
+++ b/source/include/Ostap/MoreMath.h
@@ -1459,6 +1459,19 @@ namespace Ostap
     // ========================================================================
 
     
+    // ========================================================================
+    // Lambert W-function
+    // ========================================================================
+    /** get the Lambert \f$ W_0 \f$  function for \f$- \frac{1}{e} < x \f$ 
+     *  @see https://en.wikipedia.org/wiki/Lambert_W_function
+     */
+    double lambert_W0 ( const double x ) ;
+    // ========================================================================
+    /** get the Lambert \f$ W_{-1}\f$  function for \f$- \frac{1}{e} < x < 0  \f$ 
+     *  @see https://en.wikipedia.org/wiki/Lambert_W_function
+     */
+    double lambert_Wm1 ( const double x ) ;
+    
     // ========================================================================
     /** complete Fermi-Dirac integral 
      *  \f$ F_j(x) = \frac{1}{\Gamma(j+1)}\int^{+\infty}_{0} \frac{t^j}{ \exp ( t-x) + 1 } dt \f$ 
diff --git a/source/include/Ostap/ValueWithError.h b/source/include/Ostap/ValueWithError.h
index 672ee277..454c3472 100644
--- a/source/include/Ostap/ValueWithError.h
+++ b/source/include/Ostap/ValueWithError.h
@@ -1214,6 +1214,11 @@ namespace Ostap
     ( const double          x     , 
       const ValueWithError& y     ) ;
     // ========================================================================
+    /* get the Lambert W_0 function for \f$- \frac{1}{e} < x \f$ 
+     *  @see https://en.wikipedia.org/wiki/Lambert_W_function
+     */
+    ValueWithError lambert_W0 ( const ValueWithError& x ) ;
+    // ========================================================================
     /// check for NaN
     inline bool isnan    ( const ValueWithError& v ) { return v.isnan    () ; }
     /// finite ?
diff --git a/source/src/MoreMath.cpp b/source/src/MoreMath.cpp
index fa0f10f0..f2e66137 100644
--- a/source/src/MoreMath.cpp
+++ b/source/src/MoreMath.cpp
@@ -23,6 +23,7 @@
 #include "gsl/gsl_sf_psi.h"
 #include "gsl/gsl_sf_bessel.h"
 #include "gsl/gsl_sf_ellint.h"
+#include "gsl/gsl_sf_lambert.h"
 // ============================================================================
 // LHCbMath
 // ============================================================================
@@ -3922,6 +3923,48 @@ double Ostap::Math::der_Bi ( const double x )
 }
 // ============================================================================
 
+// ============================================================================
+/* get the Lambert W_0 function for \f$- \frac{1}{e} < x \f$ 
+ *  @see https://en.wikipedia.org/wiki/Lambert_W_function
+ */
+// ============================================================================
+double Ostap::Math::lambert_W0 ( const double x )
+{
+  //
+  static const double s_mnmn = -1 / M_E ;
+  if (x <= s_mnmn ) { return std::numeric_limits<double>::quiet_NaN(); }
+  //
+  gsl_sf_result result ;
+  const int ierror = gsl_sf_lambert_W0_e ( x , &result ) ;
+  if ( ierror ) 
+    {
+      gsl_error ( "Error from gsl_sf_lambert_W0_e" , __FILE__ , __LINE__ , ierror ) ;
+      if      ( ierror == GSL_EDOM     ) // input domain error, e.g sqrt(-1)
+        { return std::numeric_limits<double>::quiet_NaN(); }
+    }
+  return result.val ;
+}
+// ============================================================================
+/*  get the Lambert W_0 function for \f$- \frac{1}{e} < x < 0  \f$ 
+ *  @see https://en.wikipedia.org/wiki/Lambert_W_function
+ */
+// ============================================================================
+double Ostap::Math::lambert_Wm1 ( const double x )
+{
+  //
+  static const double s_mnmn = -1 / M_E ;
+  if ( x <= s_mnmn | 0 <= x ) { return std::numeric_limits<double>::quiet_NaN(); }
+  //
+  gsl_sf_result result ;
+  const int ierror = gsl_sf_lambert_Wm1_e ( x , &result ) ;
+  if ( ierror ) 
+    {
+      gsl_error ( "Error from gsl_sf_lambert_Wm1_e" , __FILE__ , __LINE__ , ierror ) ;
+      if      ( ierror == GSL_EDOM     ) // input domain error, e.g sqrt(-1)
+        { return std::numeric_limits<double>::quiet_NaN(); }
+    }
+  return result.val ;
+}
 
 // ============================================================================
 /*  complete Fermi-Dirac integral 
diff --git a/source/src/ValueWithError.cpp b/source/src/ValueWithError.cpp
index acf8876e..39796674 100644
--- a/source/src/ValueWithError.cpp
+++ b/source/src/ValueWithError.cpp
@@ -3199,6 +3199,24 @@ Ostap::Math::bessel_Knu
   return Ostap::Math::ValueWithError ( value  , derivative * derivative * x.cov2 () ) ;
 }
 // ============================================================================
+/*  get the Lambert W_0 function for \f$- \frac{1}{e} < x \f$ 
+ *  @see https://en.wikipedia.org/wiki/Lambert_W_function
+ */
+// ============================================================================
+Ostap::Math::ValueWithError
+Ostap::Math::lambert_W0
+( const Ostap::Math::ValueWithError& x ) 
+{
+  const double z  = x.value () ;
+  const double c2 = x.cov2  () ;
+  const double w  = lambert_W0 ( z ) ;
+  if ( c2 <= 0 || M_E * z <= -1 ) { return w ; } 
+  //
+  const double d = s_zero ( z ) ? 1.0 : w / ( z * ( 1 + w ) ) ;
+  //
+  return Ostap::Math::ValueWithError ( x , c2 * d * d ) ; 
+}
+// ============================================================================
 /*  arithmetic-geometric mean 
  *  @param x the first value 
  *  @param y the second value