Lambertw (#24)

Internal implementation of the Lambert-W function

Co-Authored-By: [email protected]

Author: hunterboerner

lib/nx_signal.ex

@@ -181,7 +181,7 @@ defmodule NxSignal do
 
       iex> NxSignal.as_windowed(Nx.tensor([0, 1, 2, 3, 4, 10, 11, 12]), window_length: 4)
       #Nx.Tensor<
-        s64[5][4]
+        s32[5][4]
         [
           [0, 1, 2, 3],
           [1, 2, 3, 4],
@@ -193,7 +193,7 @@ defmodule NxSignal do
 
       iex> NxSignal.as_windowed(Nx.tensor([0, 1, 2, 3, 4, 10, 11, 12]), window_length: 3)
       #Nx.Tensor<
-        s64[6][3]
+        s32[6][3]
         [
           [0, 1, 2],
           [1, 2, 3],
@@ -206,7 +206,7 @@ defmodule NxSignal do
 
       iex> NxSignal.as_windowed(Nx.tensor([0, 1, 2, 3, 4, 10, 11]), window_length: 2, stride: 2, padding: [{0, 3}])
       #Nx.Tensor<
-        s64[5][2]
+        s32[5][2]
         [
           [0, 1],
           [2, 3],
@@ -219,7 +219,7 @@ defmodule NxSignal do
       iex> t = Nx.iota({7});
       iex> NxSignal.as_windowed(t, window_length: 6, padding: :reflect, stride: 1)
       #Nx.Tensor<
-        s64[8][6]
+        s32[8][6]
         [
           [3, 2, 1, 0, 1, 2],
           [2, 1, 0, 1, 2, 3],
@@ -234,7 +234,7 @@ defmodule NxSignal do
 
       iex> NxSignal.as_windowed(Nx.iota({10}), window_length: 6, padding: :reflect, stride: 2)
       #Nx.Tensor<
-        s64[6][6]
+        s32[6][6]
         [
           [3, 2, 1, 0, 1, 2],
           [1, 0, 1, 2, 3, 4],
@@ -549,7 +549,7 @@ defmodule NxSignal do
       iex> result = NxSignal.istft(z, w, opts)
       iex> Nx.as_type(result, Nx.type(t))
       #Nx.Tensor<
-        s64[8]
+        s32[8]
         [0, 10, 1, 0, 10, 10, 2, 20]
       >
 
@@ -563,7 +563,7 @@ defmodule NxSignal do
       iex> result = NxSignal.istft(z, w, opts)
       iex> Nx.as_type(result, Nx.type(t))
       #Nx.Tensor<
-        s64[8]
+        s32[8]
         [0, 10, 1, 0, 10, 10, 2, 20]
       >
 
@@ -655,21 +655,21 @@ defmodule NxSignal do
 
       iex> NxSignal.overlap_and_add(Nx.iota({3, 4}), overlap_length: 0)
       #Nx.Tensor<
-        s64[12]
+        s32[12]
         [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
       >
 
       iex> NxSignal.overlap_and_add(Nx.iota({3, 4}), overlap_length: 3)
       #Nx.Tensor<
-        s64[6]
+        s32[6]
         [0, 5, 15, 18, 17, 11]
       >
 
       iex> t = Nx.tensor([[[[0, 1, 2, 3], [4, 5, 6, 7]]], [[[10, 11, 12, 13], [14, 15, 16, 17]]]]) |> Nx.vectorize(x: 2, y: 1)
       iex> NxSignal.overlap_and_add(t, overlap_length: 3)
       #Nx.Tensor<
         vectorized[x: 2][y: 1]
-        s64[5]
+        s32[5]
         [
           [
             [0, 5, 7, 9, 7]

lib/nx_signal/internal.ex

@@ -0,0 +1,140 @@
+defmodule NxSignal.Internal do
+  @moduledoc false
+  import Nx.Defn
+
+  @omega 0.56714329040978387299997
+  @expn1 0.36787944117144232159553
+
+  deftransform lambert_w(z, k, opts \\ []) do
+    opts = Keyword.validate!(opts, tol: 1.0e-8)
+
+    z =
+      if Nx.Type.complex?(Nx.type(z)) do
+        Nx.as_type(z, :c128)
+      else
+        Nx.complex(Nx.as_type(z, :f64), 0)
+      end
+
+    lambert_w_n(z, k, tol: opts[:tol])
+  end
+
+  defnp lambert_w_n(z, k, opts) do
+    tol = opts[:tol]
+
+    rz = Nx.real(z)
+
+    cond do
+      Nx.is_infinity(rz) and rz > 0 ->
+        z + 2.0 * Nx.Constants.pi(:f64) * k * Nx.Constants.i()
+
+      Nx.is_infinity(rz) and rz < 0 ->
+        -z + 2.0 * Nx.Constants.pi(:f64) * k * Nx.Constants.i()
+
+      z == 0 and k == 0 ->
+        z
+
+      z == 0 ->
+        Nx.Constants.neg_infinity(:f64)
+
+      Nx.equal(z, 1) and k == 0 ->
+        @omega
+
+      true ->
+        halleys_method(z, k, tol)
+    end
+  end
+
+  defnp halleys_method(z, k, tol) do
+    absz = Nx.abs(z)
+
+    w =
+      cond do
+        k == 0 ->
+          cond do
+            Nx.abs(z + @expn1) < 0.3 ->
+              lambertw_branchpt(z)
+
+            -1.0 < Nx.real(z) and Nx.real(z) < 1.5 and Nx.abs(Nx.imag(z)) < 1.0 and
+                -2.5 * Nx.abs(Nx.imag(z)) - 0.2 < Nx.real(z) ->
+              lambertw_pade0(z)
+
+            true ->
+              lambertw_asy(z, k)
+          end
+
+        k == -1 and absz <= @expn1 and Nx.imag(z) == 0.0 and Nx.real(z) < 0.0 ->
+          Nx.log(-Nx.real(z))
+
+        k == -1 ->
+          lambertw_asy(z, k)
+
+        true ->
+          lambertw_asy(z, k)
+      end
+
+    # Halley's Method
+    cond do
+      Nx.real(w) >= 0 ->
+        {w, _} =
+          while {w, {z, tol, i = 0}}, i < 100 do
+            ew = Nx.exp(-w)
+            wewz = w - z * ew
+            wn = w - wewz / (w + 1.0 - (w + 2.0) * wewz / (2.0 * w + 2.0))
+
+            if Nx.abs(wn - w) <= tol * Nx.abs(wn) do
+              {wn, {z, tol, 100}}
+            else
+              {wn, {z, tol, i + 1}}
+            end
+          end
+
+        w
+
+      true ->
+        {w, _} =
+          while {w, {z, tol, i = 0}}, i < 100 do
+            ew = Nx.exp(w)
+            wew = w * ew
+            wewz = wew - z
+            wn = w - wewz / (wew + ew - (w + 2.0) * wewz / (2.0 * w + 2.0))
+
+            if Nx.abs(wn - w) <= tol * Nx.abs(wn) do
+              {wn, {z, tol, 100}}
+            else
+              {wn, {z, tol, i + 1}}
+            end
+          end
+
+        w
+    end
+  end
+
+  defnp lambertw_branchpt(z) do
+    m_e =
+      Nx.Constants.e(:f64)
+
+    p = Nx.sqrt(2.0 * (m_e * z + 1.0))
+
+    cevalpoly_2(p, -1.0 / 3.0, 1.0, -1.0)
+  end
+
+  defnp lambertw_pade0(z) do
+    z * cevalpoly_2(z, 12.85106382978723404255, 12.34042553191489361902, 1.0) /
+      cevalpoly_2(z, 32.53191489361702127660, 14.34042553191489361702, 1.0)
+  end
+
+  defnp lambertw_asy(z, k) do
+    w =
+      Nx.log(z) + 2.0 * Nx.Constants.pi(:f64) * k * Nx.Constants.i()
+
+    w - Nx.log(w)
+  end
+
+  defnp cevalpoly_2(z, c0, c1, c2) do
+    s = Nx.abs(z) ** 2
+    r = 2 * Nx.real(z)
+    b = -s * c0 + c2
+    a = r * c0 + c1
+    z * a + b
+  end
+end